
    9=eJ@                        d Z dZg dZddlZddlmZmZ ddlm	Z	m
Z
mZmZ ddlmZ dd	lmZmZmZmZmZmZ dd
lmZ  G d de          Zd Z G d dee          Zej         e_          G d dee          Z eej                   e_         dS )zSparse DIAgonal formatzrestructuredtext en)	dia_array
dia_matrixisspmatrix_dia    N   )spmatrix_array_doc_to_matrix)issparse_formats_spbasesparray)_data_matrix)isshapeupcast_chargetdtypeget_sum_dtypevalidateaxischeck_shape)
dia_matvecc                      e Zd ZdZdZddZd Zd Zd Zdd	Z	e
j	        j        e	_        e
j        j        e_        dd
Ze
j        j        e_        d Zd Zd ZddZddZe
j        j        e_        ddZe
j        j        e_        ddZe
j        j        e_        ddZe
j        j        e_        ddZe
j        j        e_        ddZd Ze
j        j        e_        dS )	_dia_basea  Sparse matrix with DIAgonal storage

    This can be instantiated in several ways:
        dia_array(D)
            with a dense matrix

        dia_array(S)
            with another sparse matrix S (equivalent to S.todia())

        dia_array((M, N), [dtype])
            to construct an empty matrix with shape (M, N),
            dtype is optional, defaulting to dtype='d'.

        dia_array((data, offsets), shape=(M, N))
            where the ``data[k,:]`` stores the diagonal entries for
            diagonal ``offsets[k]`` (See example below)

    Attributes
    ----------
    dtype : dtype
        Data type of the matrix
    shape : 2-tuple
        Shape of the matrix
    ndim : int
        Number of dimensions (this is always 2)
    nnz
        Number of stored values, including explicit zeros
    data
        DIA format data array of the matrix
    offsets
        DIA format offset array of the matrix

    Notes
    -----

    Sparse matrices can be used in arithmetic operations: they support
    addition, subtraction, multiplication, division, and matrix power.

    Examples
    --------

    >>> import numpy as np
    >>> from scipy.sparse import dia_array
    >>> dia_array((3, 4), dtype=np.int8).toarray()
    array([[0, 0, 0, 0],
           [0, 0, 0, 0],
           [0, 0, 0, 0]], dtype=int8)

    >>> data = np.array([[1, 2, 3, 4]]).repeat(3, axis=0)
    >>> offsets = np.array([0, -1, 2])
    >>> dia_array((data, offsets), shape=(4, 4)).toarray()
    array([[1, 0, 3, 0],
           [1, 2, 0, 4],
           [0, 2, 3, 0],
           [0, 0, 3, 4]])

    >>> from scipy.sparse import dia_array
    >>> n = 10
    >>> ex = np.ones(n)
    >>> data = np.array([ex, 2 * ex, ex])
    >>> offsets = np.array([-1, 0, 1])
    >>> dia_array((data, offsets), shape=(n, n)).toarray()
    array([[2., 1., 0., ..., 0., 0., 0.],
           [1., 2., 1., ..., 0., 0., 0.],
           [0., 1., 2., ..., 0., 0., 0.],
           ...,
           [0., 0., 0., ..., 2., 1., 0.],
           [0., 0., 0., ..., 1., 2., 1.],
           [0., 0., 0., ..., 0., 1., 2.]])
    diaNFc           
         t          j        |            t          |          r|j        dk    rI|r|                                }|j        | _        |j        | _        t          |j                  | _	        n|j        | j        k    r|r|                                }n|
                                }|j        | _        |j        | _        t          |j                  | _	        nt          |t                    rmt          |          rt          |          | _	        t          j        dt!          |t"                              | _        |                     t'          | j                            }t          j        d|          | _        nr	 |\  }}|t)          d          t          j        t          j        |d         ||                    | _        t          j        t          j        |d	         |                     t'          |                    |                    | _        t          |          | _	        n# t0          $ r}	t)          d
          |	d }	~	ww xY w	 t          j        |          }n*# t0          $ r}	t)          d| j        z            |	d }	~	ww xY w|                     |||          
                                }|j        | _        |j        | _        t          |j                  | _	        || j                            |          | _        | j        j        d	k    rt)          d          | j        j        dk    rt)          d          | j        j        d         t;          | j                  k    r6t)          d| j        j        d         t;          | j                  fz            t;          t          j        | j                            t;          | j                  k    rt)          d          d S )Nr   )r   r   )defaultmaxvalr   dtypezexpected a shape argument)r   copyr   z+unrecognized form for dia_array constructorz+unrecognized form for %s_matrix constructor)r   shapezoffsets array must have rank 1   zdata array must have rank 2zBnumber of diagonals (%d) does not match the number of offsets (%d)z&offset array contains duplicate values)r   __init__r	   formatr   dataoffsetsr   r   _shapetodia
isinstancetupler   npzerosr   float_get_index_dtypemax
ValueError
atleast_2darray
atleast_1d	Exceptionasarray_coo_containerastypendimlenunique)
selfarg1r   r   r   A	idx_dtyper#   r$   es
             1lib/python3.11/site-packages/scipy/sparse/_dia.pyr!   z_dia_base.__init__Y   s   d###D>> /	/{e## '99;;D I	#|)$*55;$+--$-		AA

AF	 y)!'22e$$  	/t}} 5 *$//HUHUE,J,J,JKK	 11TZ1II	!x9===5$(MD' }()DEEE "bhtAweRV.W.W.W X XDI#%=$q'@D@U@U]`af]g]g@U@h@h?C2E 2E 2E $F $FDL #.e"4"4DKK ! [ [ [$%RSSYZZ[Gz$'' G G G  "137;"? @ @EFGG ##DU#CCIIKKADI9DL%ag..DK	((//DI <!!=>>>9>Q:;;;9?1T\!2!222 @yq)3t|+<+<=> ? ? ? ry&&''3t|+<+<<<EFFF =<s0   I 
I&I!!I&*I? ?
J&	J!!J&c                     t           |                                          d         }d| j        | j        j        | j        | j        j        d         |fz   z  S )Nr   zW<%dx%d sparse matrix of type '%s'
	with %d stored elements (%d diagonals) in %s format>r   )r
   	getformatr   r   typennzr#   )r9   r"   s     r>   __repr__z_dia_base.__repr__   sR    $..**+A.H
djotx9K$& &' 	'    c                     | j         \  }}t          j        | j        j         d                   }|| j        dddf         z
  }|dk    }|||k     z  }|||k     z  }|S )z~Returns a mask of the same shape as self.data, where
        mask[i,j] is True when data[i,j] corresponds to a stored element.r   Nr   )r   r)   aranger#   r$   )r9   num_rowsnum_colsoffset_indsrowmasks         r>   
_data_maskz_dia_base._data_mask   sl     "Z(i	 233DL400qx x'(rD   c                 h    |                                  }t          j        | j        |                   S N)rL   r)   count_nonzeror#   )r9   rK   s     r>   rO   z_dia_base.count_nonzero   s(      	$000rD   c                     |t          d          | j        \  }}d}| j        D ]5}|dk    r|t          |||z
            z  }|t          ||z   |          z  }6t	          |          S )Nz6_getnnz over an axis is not implemented for DIA formatr   )NotImplementedErrorr   r$   minint)r9   axisMNrB   ks         r>   _getnnzz_dia_base._getnnz   s    % '7 8 8 8j! 	" 	"A1uus1QqSzz!s1Q3qzz!3xxrD   c           
         t          |           ||dk     r|dz  }t          | j                  }| j        \  }}d }|dk    r|                                 }| j        |z                      d          }	|	j        d         |k    r|	}
n-t          j        ||	j                  }
|	|
d |	j        d         <   | 	                    |
|          }nt          j        |df|          }t          j
        ||          }t          ||t          | j                  | j        j        d         | j        | j        ||           | 	                    |          }||                    ||          S | 	                    |                    |                    }||j        |j        k    rt          d          |                    d||	          S )
Nr   r    rT   r   r   )r   outzdimensions do not match )rT   r   r[   )r   r   r   r   rL   r#   sumr)   r*   _ascontaineronesr   r7   r$   r.   )r9   rT   r   r[   	res_dtyperG   rH   retrK   xresrow_sumsones                r>   r]   z_dia_base.sum   s   TqAID!$*--	!Z(199??$$DT!&&A&..AwqzX%%hxqw777#$KQWQZK ##Cy#99CC x1Y???H'()444Cx3t|+<+<yq)4<CS S S ((22H|||%S|999##HLLdL$;$;<<C?syCI556777wwBew555rD   c                    t          |t          |                     s'|                                                     |          S t	          j        | j        |j                  }|                     |d          }|j        D ]t}||j        v r@|	                    |
                    |          |
                    |          z   |           K|	                    |
                    |          |           u|S )NTr   )r'   rA   tocsr_add_sparser)   promote_typesr   r5   r$   setdiagdiagonal)r9   otherr   mds        r>   ri   z_dia_base._add_sparse   s     %d,, 	3 ::<<++E222
  U[99KKDK))  	0 	0AAI~~ 		!**Q--%..*;*;;Q????		%..++Q////rD   c                 t   |}t          j        | j        d         t          | j        j        |j        j                            }| j        j        d         }| j        \  }}t          ||t          | j	                  || j	        | j        |
                                |
                                           |S )Nr   r   r   )r)   r*   r   r   r   charr#   r   r7   r$   ravel)r9   rm   rb   yLrU   rV   s          r>   _mul_vectorz_dia_base._mul_vector  s    HTZ]+djo78w|+E +E F F F IOAj!1QDL))1dlDIqwwyyRSRYRYR[R[\\\rD   c                 N     t          j         fd|j        D                       S )Nc                 b    g | ]+}                     |                              d d          ,S )r   )ru   reshape).0colr9   s     r>   
<listcomp>z._dia_base._mul_multimatrix.<locals>.<listcomp>  s7    QQQ#$**3//771==QQQrD   )r)   hstackT)r9   rm   s   ` r>   _mul_multimatrixz_dia_base._mul_multimatrix  s*    yQQQQQQQRRRrD   r   c                    | j         \  }}|j        dk    rt          j        }nt	          |          }|dk     rt          ||z   ||          }d}|}nt          |||z
  |          }|}||z   }|j        dk    r
|d |         }| j        j         \  }	}
|| j        v rX||
k    r9t          j        |	|f| j        j	                  }| j        |d d d |
f<   || _        || j        | j        |k    ||f<   d S t          j
        | j        | j        j	                            |                    | _        t          ||
          }t          j        |	dz   |f| j        j	                  }| j        |d dd |
f<   ||d||f<   || _        d S )Nr   r   r   rx   )r   r6   r)   infr7   rR   r#   r$   r*   r   appendrA   r-   )r9   valuesrW   rU   rV   values_nn	min_index	max_index	data_rows	data_colsr#   rn   s                r>   _setdiagz_dia_base._setdiag  s   z1;!vHH6{{Hq55AE1h''AIIIAq1uh''AIAI;!BQBZF#y	99$$xI 6dioNNN&*iQQQ

]# 	@FDIdla'9)<<===9T\4<3E3J3J13M3MNNDLIy))A8Y]A.dioFFFD$(ID"jyj!,2DYy(()DIIIrD   c                 2    |r|                                  S | S rN   rg   )r9   r   s     r>   r&   z_dia_base.todia=  s     	99;;KrD   c                 f   |t          d          | j        \  }}t          | j                  }| j         }t	          j        t          |          t          j                  d d d f         }t	          j        |t          j                  ||z  d d d f         z
  }t          d|| j        j        d         z
            }	t	          j	        | j        t	          j
        | j        j        d         |	f| j        j                  f          }
|
||f         }
|                     |
|f||f|          S )NzoSparse matrices do not support an 'axes' parameter because swapping dimensions is the only logical permutation.r   r   r   )r   r   )r.   r   r-   r$   r)   rF   r7   intcr#   r}   r*   r   _dia_container)r9   axesr   rG   rH   max_dimr$   rc
pad_amountr#   s              r>   	transposez_dia_base.transposeE  sE    L M M M "Z(dj// <- Ic'll"'222111d7;Ihbg...'G2CQQQW1MMGDIOA$6677
y$)RXtyq/A:.N48IO&E &E &E F G GAqDz""D'?h; &* # , , 	,rD   c                    | j         \  }}|| k    s||k    r t          j        d| j        j                  S t          j        | j        |k              \  }t          d|          }t          ||z   |          }||z
  }|j	        dk    r t          j
        || j        j                  S | j        |d         ||f         }|t          |          z
  }	|	dk    rt          j        |d|	fd          }|S )Nr   r   constant)mode)r   r)   emptyr#   r   nonzeror$   r-   rR   sizer*   r7   pad)
r9   rW   rowscolsidx	first_collast_colresult_sizeresultpaddings
             r>   rl   z_dia_base.diagonal]  s    Z
d::d8ATY_5555z$,!+,,1II	tax&&*8q==8Kty????3q69X#556F+Q;;VFQLzBBBFrD   c                    | j         dk    r!|                     | j        | j                  S | j        \  }}| j        j        \  }}t          j        |          }|| j        d d d f         z
  }|dk    }|||k     z  }|||k     z  }|| j        dk    z  }|                     t          | j                            }	t          j
        |dz   |	          }
t          j        |                    d          d |                   |
d|dz   <   ||k     r|
|         |
|dz   d <   |j        |j                                     |	d          }| j        j        |j                 }|                     |||
f| j        | j                  S )	Nr   r   r   r   rZ   Frg   )r   r   )rB   _csc_containerr   r   r#   r)   rF   r$   r,   r-   r*   cumsumr]   r~   r5   )r9   r   rG   rH   num_offsets
offset_lenrI   rJ   rK   r<   indptrindicesr#   s                r>   tocscz_dia_base.tocsco  s   8q==&&tz&DDD!Z("&)/Zi
++DL400qx x'(a ))TZ)AA	(Q,i888!#4888+;+;IXI+F!G!GqA~  $*:$6F:a<==!%-&&yu&==y{46"""D'6#:$*)- # 5 5 	5rD   c                    | j         \  }}| j        j         \  }}t          j        |          }|| j        d d d f         z
  }|dk    }|||k     z  }|||k     z  }|| j        dk    z  }||         }t          j        ||          |                                         }	| j        |         }
|                     |
||	ff| j         | j        d          S )Nr   F)r   r   r   )	r   r#   r)   rF   r$   tilerr   r4   r   )r9   r   rG   rH   r   r   rI   rJ   rK   r{   r#   s              r>   tocooz_dia_base.tocoo  s    !Z("&)/Zi
++DL400qx x'(a $igk;//

=y ""C:dj
 # 
 
 	
rD   Tc                     |r5|                      || j                                        f| j                  S |                      || j        f| j                  S )zReturns a matrix with the same sparsity structure as self,
        but with different data.  By default the structure arrays are copied.
        )r   )r   r$   r   r   )r9   r#   r   s      r>   
_with_dataz_dia_base._with_data  sp      	&&t|((**+4: '    &&t|$DJ '   rD   c                    t          |          }|\  }}| j        d d d |f         | _        || j        d         k    rt          j        | j        | j        d         z   | j        j        d         k               rO| j        d d d f         | j        d         z   t          j        | j        j        d                   k    }d| j        |<   || _        d S )Nr   r   )r   r#   r   r)   anyr$   rF   r%   )r9   r   rU   rV   rK   s        r>   resizez_dia_base.resize  s    E""1Iaaa!e$	
1t|djm3dioa6HHII  LD)DJqM9Idioa0112DDIdOrD   )NNFrN   )NNN)r   )F)NF)T)__name__
__module____qualname____doc___formatr!   rC   rL   rO   rX   r   r]   ri   ru   r   r   r&   r   rl   r   r   r   r   r\   rD   r>   r   r      s       E EL GDG DG DG DGL' ' '	 	 	1 1 1    o-GO#19M$6 $6 $6 $6L +%CK  4  S S S# # # #J    M)EM, , , ,,  )1I     '/H5 5 5 50 M)EM
 
 
 
& M)EM      ^+FNNNrD   r   c                 ,    t          | t                    S )a  Is `x` of dia_matrix type?

    Parameters
    ----------
    x
        object to check for being a dia matrix

    Returns
    -------
    bool
        True if `x` is a dia matrix, False otherwise

    Examples
    --------
    >>> from scipy.sparse import dia_array, dia_matrix, coo_matrix, isspmatrix_dia
    >>> isspmatrix_dia(dia_matrix([[5]]))
    True
    >>> isspmatrix_dia(dia_array([[5]]))
    False
    >>> isspmatrix_dia(coo_matrix([[5]]))
    False
    )r'   r   )rb   s    r>   r   r     s    . a$$$rD   c                       e Zd ZdS )r   Nr   r   r   r\   rD   r>   r   r             DrD   r   c                       e Zd ZdS )r   Nr   r\   rD   r>   r   r     r   rD   r   )r   __docformat____all__numpyr)   _matrixr   r   _baser	   r
   r   r   _datar   _sputilsr   r   r   r   r   r   _sparsetoolsr   r   r   r   r   r\   rD   r>   <module>r      s|    %
7
7
7     3 3 3 3 3 3 3 3 7 7 7 7 7 7 7 7 7 7 7 7       ` ` ` ` ` ` ` ` ` ` ` ` ` ` ` ` $ $ $ $ $ $k, k, k, k, k, k, k, k,\% % %6	 	 	 	 		7 	 	 	 %	 	 	 	 	 	9 	 	 	 *))*;<<
   rD   