a jg5@sdZddlmZddlmZddlZddlmZer@ddl mZ e ej fZ ddZGd d d Zd d d dZddZddZdddZddZddZdS)z0Indexing mixin for sparse array/matrix classes. ) annotations) TYPE_CHECKINGN) isintlikecCs0t||\}}|jj|j_|jj|j_||fS)a  Same as np.broadcast_arrays(a, b) but old writeability rules. NumPy >= 1.17.0 transitions broadcast_arrays to return read-only arrays. Set writeability explicitly to avoid warnings. Retain the old writeability rules, as our Cython code assumes the old behavior. )npZbroadcast_arraysflagsZ writeable)abxyr 9E:\soccer\ml_env\Lib\site-packages\scipy/sparse/_index.py_broadcast_arrayss   rc@seZdZdZddZddZddZdd Zd d Zd d Z ddZ ddZ ddZ ddZ ddZddZddZddZddZd d!Zd"d#Zd$d%Zd&d'Zd(d)Zd*S)+ IndexMixinzS This class provides common dispatching and validation logic for indexing. cCs"ddlm}t||rtddS)zWe do not currently support 1D sparse arrays. This function is called each time that a 1D array would result, raising an error instead. Once 1D sparse arrays are implemented, it should be removed. r)sparrayzcWe have not yet implemented 1D sparse slices; please index using explicit indices, e.g. `x[:, [0]]`N)Z scipy.sparser isinstanceNotImplementedError)selfrr r r _raise_on_1d_array_slice#s   z#IndexMixin._raise_on_1d_array_slicecCs||\}}t|trt|tr.|||St|trL||||S|jdkrj||||S|jdkr|||St dn:t|tr t|tr|| ||St|tr|tdkr||kr| S| ||S|jdkr| ||St dn|jdkrPt|tr6||||St|tr|||Snvt|trh|||St|tr~t dnH|jddkr|jdks|jddkr||dddf|St||\}}|j|jkrt d|jdkr|jt|j|jdS|||S)Nrzindex results in >2 dimensionsr'number of row and column indices differdtype)_validate_indicesr INT_TYPES _get_intXintslicer_get_intXslicendim_get_intXarray IndexError_get_sliceXintcopy_get_sliceXslice_get_sliceXarray_get_arrayXint_get_arrayXsliceshape_get_columnXarrayZravelrsize __class__rZ atleast_2dr_get_arrayXarray)rkeyrowcolr r r __getitem__3sT                          , zIndexMixin.__getitem__c Csb||\}}t|tr\t|tr\tj||jd}|jdkrDtd||||j ddSt|t rtj | |j ddddf}n t|}t|t rtj | |j ddddf}|jdkr|dddf}n t|}t||\}}|j |j kr tdddlm}||r|jdkr<|d}|d}|j ddkoX|j ddk}|j ddkov|j ddk} |s|j d|j dkr| s|j d|j dkstd|j ddks|j ddkrdS|jdd }|||||n^tj||jd}|j |j kr4t||j }|jdkrDdS||j }||||dS) Nrrz&Trying to assign a sequence to an itemrrissparsezshape mismatch in assignmentT)r")rrrrasarrayrr) ValueError _set_intXintZflatrZarangeindicesr'Z atleast_1drrr _baser1ZtocooZsum_duplicates_set_arrayXarray_sparseZsqueezeZ broadcast_toZreshape_set_arrayXarray) rr,r r-r.ijr1Z broadcast_rowZ broadcast_colr r r __setitem__esV  $  "         zIndexMixin.__setitem__c Cszddlm}t||tjfrT|jdkrT|jjdkrT|j|jkrFt d| \}}n t |\}}|j\}}dddd d d d }t |rt |}|| ks||krt d ||dkr||7}n4t|}dur|||d}nt|ts|||}t |r:t |}|| ks||kr&t d||dkrr||7}n8t|} durZ|| |d}nt|tsr|||}||fS)Nr)_spbaserr z.boolean index shape does not match array shapeznpt.NDArray[np.bool_]intstrznpt.NDArray[np.int_])idx axis_size axis_namereturncSs2t||kr*td|dt|d|t|S)Nzboolean z index has incorrect length: z instead of )lenr _boolean_index_to_array)r?r@rAr r r _validate_bool_idxs z8IndexMixin._validate_indices.._validate_bool_idxzrow index (%d) out of rangerr-zcolumn index (%d) out of rangecolumn)r6r<rrZndarrayrrkindr'r nonzero _unpack_indexrr=_compatible_boolean_indexr _asindices) rr,r<r-r.MNrEZbool_rowZbool_colr r r rsB                zIndexMixin._validate_indicesc Cszt|}Wn4tttfyB}ztd|WYd}~n d}~00|jdvrVtd|jdkrd|S|}||krtd|| }|dkr|| krtd|||us|j j s| }||dk|7<|S)zConvert `idx` to a valid index for an axis with a given length. Subclasses that need special validation can override this method. z invalid indexN)rrzIndex dimension must be 1 or 2rindex (%d) out of range) rr2r3 TypeError MemoryErrorr rr)maxminrZowndatar")rr?lengthr eZmax_indxZmin_indxr r r rKs&      zIndexMixin._asindicescCsP|j\}}t|}|| ks$||kr0td||dkr@||7}||tdS)zGReturn a copy of row i of the matrix, as a (1 x n) row vector. rNrN)r'r=r rrrr9rLrMr r r _getrows  zIndexMixin._getrowcCsP|j\}}t|}|| ks$||kr0td||dkr@||7}|td|S)zMReturn a copy of column i of the matrix, as a (m x 1) column vector. rNrN)r'r=r r!rrUr r r _getcols  zIndexMixin._getcolcCs tdSNrrr-r.r r r rszIndexMixin._get_intXintcCs tdSrXrYrZr r r rszIndexMixin._get_intXarraycCs tdSrXrYrZr r r rszIndexMixin._get_intXslicecCs tdSrXrYrZr r r r!szIndexMixin._get_sliceXintcCs tdSrXrYrZr r r r# szIndexMixin._get_sliceXslicecCs tdSrXrYrZr r r r$ szIndexMixin._get_sliceXarraycCs tdSrXrYrZr r r r%szIndexMixin._get_arrayXintcCs tdSrXrYrZr r r r&szIndexMixin._get_arrayXslicecCs tdSrXrYrZr r r r(szIndexMixin._get_columnXarraycCs tdSrXrYrZr r r r+szIndexMixin._get_arrayXarraycCs tdSrXrYrr-r.r r r r r4szIndexMixin._set_intXintcCs tdSrXrYr[r r r r8szIndexMixin._set_arrayXarraycCs4tj||jd}t||\}}||||dS)Nr)rr2Ztoarrayrrr8)rr-r.r _r r r r7!sz"IndexMixin._set_arrayXarray_sparseN)__name__ __module__ __qualname____doc__rr/r;rrKrVrWrrrr!r#r$r%r&r(r+r4r8r7r r r r rs*240  rzctuple[int | slice | npt.NDArray[np.bool_ | np.int_], int | slice | npt.NDArray[np.bool_ | np.int_]])rBcCst|}t|trRt|dkr(|\}}qt|dkrH|dtd}}qtdnHt|}|durr|td}}n(|jdkr|tdfS|jdkr|Sddl m }||s||rtd||fS)z Parse index. Always return a tuple of the form (row, col). Valid type for row/col is integer, slice, array of bool, or array of integers. rrrNzinvalid number of indicesr0zoIndexing with sparse matrices is not supported except boolean indexing where matrix and index are equal shapes.) _check_ellipsisrtuplerCrr rJrrHr6r1)indexr-r.r?r1r r r rI(s(         rIcCs|turtdtdfSt|ts&|Sddt|D}|s@|St|dkrTtd|\}tddt|}|d|tdf|||ddS)z6Process indices with Ellipsis. Returns modified index.NcSsg|]\}}|tur|qSr )Ellipsis).0r9vr r r Wz#_check_ellipsis..rz0an index can only have a single ellipsis ('...')r)rdrrrb enumeraterCr rQ)rcZellipsis_indicesr9Z num_slicesr r r raNs  racCst|}|jjdkr|SdS)z8Returns a compatible array if elements are boolean. r N)rZ asanyarrayrrGr?r r r _maybe_bool_ndarraycs  rlrcCsP|dkr dSztt|d}Wnty2YdS0t|trBdSt||dS)zQReturns True if first element of the incompatible array type is boolean. rNT)nextiterrOrbool_first_element_bool)r?Zmax_dimfirstr r r rpls  rpcCst|dst|rt|SdS)ztReturns a boolean index array that can be converted to integer array. Returns None if no such array exists. rN)hasattrrprlrkr r r rJ{srJcCs |jdkrtdt|dS)Nrzinvalid index shaper)rr rwhererkr r r rDs rD)r)r` __future__rtypingrnumpyrZ_sputilsrZ numpy.typingZnptr=integerrrrrIrarlrprJrDr r r r s"      &