:vj\ldZddlZddlZddlmZddlmZddlmZ ddl m Z d ZGd d eZGd deZGddeZd"dZGddeZGddeZGddeZGddeZGddeZGddeZGddeZd Zd!Zy#e $rZ dd lmZee Z YdZ [ dZ [ wwxYw)#z7OO interfaces to encodings for ND arrays which caching.N)caching)ABC) runlength)sparse)ExceptionWrapperct|}ttjd|zt}tj|dft }||dddf<||fS)Nrshapedtyperr)len DenseEncodingnpzerosboolint)r num_dimsencodingpaddings g/workspace/moltex-experiments/.pixi/envs/default/lib/python3.12/site-packages/trimesh/voxel/encoding.py_empty_strippedrsP5zHRXXD8O4HIHhhh]#6GGAqDM W ceZdZdZdZeejdZeejdZ eejdZ eejdZ eejdZ eejdZ eejd Zejd Zejd Zejd Zejd ZedZej*dZdZdZedZdZedZddZedZedZej@fdZ!ej@fdZ"dZ#dZ$edZ%e%jLdZ%y) Encodingaf Base class for objects that implement a specific subset of of ndarray ops. This presents a unified interface for various different ways of encoding conceptually dense arrays and to interoperate between them. Example implementations are ND sparse arrays, run length encoded arrays and dense encodings (wrappers around np.ndarrays). cp||_tj|jj|_y)N) id_function)_datarCache__hash___cache)selfdatas r__init__zEncoding.__init__&s# mm 0C0CD rcyNr#s rrzEncoding.dtype* rcyr'r(r)s rr zEncoding.shape/r*rcyr'r(r)s rsumz Encoding.sum4r*rcyr'r(r)s rsizez Encoding.size9r*rcyr'r(r)s rsparse_indiceszEncoding.sparse_indices>r*rcyr'r(r)s r sparse_valueszEncoding.sparse_valuesCr*rcyr'r(r)s rdensezEncoding.denseHr*rcyr'r(r#indicess r gather_ndzEncoding.gather_ndM rcyr'r(r#masks rr=z Encoding.maskQr:rcyr'r(r#indexs r get_valuezEncoding.get_valueUr:rcyr'r(r)s rcopyz Encoding.copyYr:rcT|j|jdk7jdk(SNr)r1r3r/r)s ris_emptyzEncoding.is_empty]s)""4#5#5#:;@@AEErcx|jrt|jS|j}|j}t |}g}g}t |D]\}}t t|t t|dz|z}tj||} tj| \} | j} | jdz} |j| || z g|jt| | t|t |tj |t"fS)z Get encoding with all zeros stripped from the start and end of each axis. Returns ------------ encoding: ? padding : (n, 2) int Padding at the start and end that was stripped raxis)rFrr r5r enumeratetuplerangeranynonzerominmaxappendslicerarrayr) r#r5r ndimsrslicesdimr/rIfilledr8loweruppers rstrippedzEncoding.strippedas =="4::. .  E "5)ICs$uU37E-B'CCDVVE-FF+JWKKMEKKMA%E NNE4%<0 1 MM%u- .*U5=12BHHWc4JJJrct||Sr')FlippedEncodingr#axess r_flipzEncoding._flip~stT**rc6|jjSz Get the hash of the current transformation matrix. Returns ------------ hash : str Hash of transformation matrix rr!r)s rr!zEncoding.__hash__zz""$$rc,t|jSr')rr r)s rrTzEncoding.ndimss4::rcNt|dk(r |jSt||SNr)rflatShapedEncodingr#r s rreshapezEncoding.reshapes"J!OtyyLe1LLrct|Sr')FlattenedEncodingr)s rrgz Encoding.flats  &&rct||Sr')_flipped)r#rIs rflipz Encoding.flipsd##rc2|j|jfSr'r1r3r)s rsparse_componentszEncoding.sparse_components""D$6$666rc|jSr'rr)s rr$z Encoding.data zzrcx|jdk7r tdtj|j|SNrz/`run_length_data` only valid for flat encodingsr)rT ValueErrorr dense_to_rler5r#rs rrun_length_datazEncoding.run_length_datas0 ::?NO O%%djj>>rcx|jdk7r tdtj|j|Srx)rTrzr dense_to_brler5r|s rbinary_run_length_datazEncoding.binary_run_length_datas0 ::?NO O&&tzz??rct||Sr') _transposedr#perms r transposezEncoding.transposes4&&rct||Sr')TransposedEncodingrs r _transposezEncoding._transposes!$--rc.|jjSr'rmutabler)s rrzEncoding.mutableszz!!!rc&||j_yr'r)r#values rrzEncoding.mutables" rNr )'__name__ __module__ __qualname____doc__r%propertyabcabstractmethodrr r-r/r1r3r5r9r=rArCrFrcache_decoratorrZr_r!rTrjrgrorrr$rint64r}rrrrsetterr(rrrrsME                          FF KK8+ %M''$77%'HH? ,.88@ '."" ^^##rrcBeZdZdZfdZedZedZejdZ ejdZ edZ edZ ejd Zejd Zd Zed Zd ZdZdZdZdZdZedZdZxZS)rz:Simple `Encoding` implementation based on a numpy ndarray.ct|tjs:t|tjs t dtj |}t|!|y)Nz(DenseEncoding data must be a numpy arrayr$) isinstancer TrackedArrayrndarrayrz tracked_arraysuperr%r#r$ __class__s rr%zDenseEncoding.__init__sM$ 4 45dBJJ/ !KLL((.D d#rc.|jjSr'rrr)s rrzDenseEncoding.dtypezzrc.|jjSr'rr r)s rr zDenseEncoding.shaperrc6|jjSr'rr-r)s rr-zDenseEncoding.sumszz~~rcBtj|j Sr'rrMrr)s rrFzDenseEncoding.is_emptys66$**%%%rc.|jjSr'rr/r)s rr/zDenseEncoding.sizezzrcD|j}|j|}||fSr')r1gather)r#r8valuess rrrzDenseEncoding.sparse_componentss%%%W%rcftjtj|jSr')r column_stackwhererr)s rr1zDenseEncoding.sparse_indicessrxx 344rc |jdSrfrrr)s rr3zDenseEncoding.sparse_values%%a((rcj|j}|D]}tj||}t|Sr')r5rror)r#r^r5as rr_zDenseEncoding._flips0 AGGE1%EU##rc|jSr'rur)s rr5zDenseEncoding.denservrc |j|Sr'rur7s rrzDenseEncoding.gatherszz'""rcF|jt|jSr')rrKTr7s rr9zDenseEncoding.gather_ndszz% *++rcp|jt|tjr|S|jSr')rrrrr5r<s rr=zDenseEncoding.masks*zz*T2::">$OODJJOOrc2|jt|Sr')rrKr?s rrAzDenseEncoding.get_valueszz%,''rcJt|jj|Sr'rrrjris rrjzDenseEncoding.reshapesTZZ//677rcJt|jj|Sr')rrrrs rrzDenseEncoding._transposesTZZ11$788rcJt|jjdSN)rr)s rrgzDenseEncoding.flatsTZZ//677rcHt|jjSr')rrrCr)s rrCzDenseEncoding.copysTZZ__.//r)rrrrr%rrr rrr-rFr/rrr1r3r_r5rr9r=rArjrrgrC __classcell__rs@rrrsD$        &&  55 ))$ #,P(89880rrcXeZdZdZdfd ZedZdZedZ edZ edZ e jdZed Zed Zed Zed Ze jd Ze jdZdZdZdZdZdZe jdZxZS)SparseEncodingz `Encoding` implementation based on an ND sparse implementation. Since the scipy.sparse implementations are for 2D arrays only, this implementation uses a single-column CSC matrix with index raveling/unraveling. c^tj}t| |||d<||d<|d}t |j dk7rt d|j |dj |j dfk7r!t dj|d|d|'t|djdd z|_ n=t||_ tj||jks t d tj|dk\s t d y) a Parameters ------------ indices: (m, n)-sized int array of indices values: (m, n)-sized dtype array of values at the specified indices shape: (n,) iterable of integers. If None, the maximum value of indices + 1 is used. r8rrzindices must be 2D, got shaped rz1values and indices shapes inconsistent: {} and {}NrHrz#all indices must be less than shapez all indices must be non-negative) r DataStorerr%rr rzformatrKrP_shaperall)r#r8rr r$rs rr%zSparseEncoding.__init__s!  " !YXy/ w}}  ">w}}>OPQ Q >  GMM!$4#6 6CJJNDO  =Y 3 3 3 ;a ?@DK,DK66'DKK/0 !FGGvvgl#?@ @$rctj|}||}ttj|d||jS)NrrHr )rrrstackr ) dense_datar1rs r from_densezSparseEncoding.from_dense3s?*-N+ HH^" -vZ=M=M  rct|jj|jj|jS)N)r8rr )rr1rCr3r r)s rrCzSparseEncoding.copy;s:'',,.%%**,**  rc |jdS)Nr8rur)s rr1zSparseEncoding.sparse_indicesBszz)$$rc |jdS)Nrrur)s rr3zSparseEncoding.sparse_valuesFszz(##rc.|jjSr')r3rr)s rrzSparseEncoding.dtypeJs!!'''rc6|jjSr')r3r-r)s rr-zSparseEncoding.sumNs!!%%''rc4|jjdSNr)r1r r)s rrTzSparseEncoding.ndimsRs""((,,rc|jSr'rr)s rr zSparseEncoding.shapeV {{rc@tj|jSr')rprodr r)s rr/zSparseEncoding.sizeZswwtzz""rc2|j|jfSr'rqr)s rrrz SparseEncoding.sparse_components^rsrc|j}tj|j|j}|j |tj ||jS)Nr )out)_cscrrr rtodenserj)r#rr5s rr5zSparseEncoding.densebsFv||6<<@5!zz%,,rc|j}|j|j}dt|g}t j |||f|j dfS)Nrrr)r3 _flat_indicesr1rsp csc_matrixr/)r#rr8indptrs rrzSparseEncoding._cscjsQ##$$T%8%89S\"}}fgv6tyy!nMMrc|jddk(rt|jdk(sJtj|j|jS)Nrr)r rrravel_multi_indexrr7s rrzSparseEncoding._flat_indicesqsD}}Q1$W]]);q)@@@##GIItzz::rchtjtj||jSr'rr unravel_indexr )r# flat_indicess r_shaped_indiceszSparseEncoding._shaped_indicesu"r// djjIJJrc|j|j|j}tj|j dSNrrH)rrrrasarraysqueeze)r#r8mats rr9zSparseEncoding.gather_ndxsBii**734<<>zz#&&B&//rctj|j|jd\}}|j |Sr)rrrrjr)r#r=i_s rr=zSparseEncoding.mask~s7xx $,,u"5671##A&&rcT|jtj|ddS)NrrH) _gather_ndr expand_dimsr?s rrAzSparseEncoding.get_values!r~~e!<=a@@rcD|jrt|jS|j}t j |d}t j |d}|dz}||jz }t j||f}t||z |j|fS)aG Get encoding with all zeros stripped from the start/end of each axis. Returns: encoding: SparseEncoding with same values but indices shifted down by padding[:, 0] padding: (n, 2) array of ints denoting padding at the start/end that was stripped rrHr) rFrr r1rrOrPrrr3)r#r8pad_left pad_rightrs rrZzSparseEncoding.strippeds =="4::. .%%66'*FF7+ R TZZ //8Y"78g0$2D2DEwNNrr')rrrrr% staticmethodrrCrr1r3rrrr-rTr r/rrr5rrrr9r=rArZrrs@rrr s:A>   %%$$(( ((--##77 -- NN ;K0 'A OOrrcjt|tj|jdft|S)ay Convenient factory constructor for SparseEncodings with values all ones. Parameters ------------ indices: (m, n) sparse indices into conceptual rank-n array shape: length n iterable or None. If None, maximum of indices along first axis + 1 is used Returns ------------ rank n bool `SparseEncoding` with True values at each index. rr )rronesr r)r8r s rSparseBinaryEncodingrs+ '277'--2B1DD#QSX YYrcBeZdZdZdfd Zej dZedZ edZ edZ dZ e ejejfdZe dd Ze dd Zej d Zej d Zej d ZdZej dZej dZej dZej dZdZdZdZdZdZdZ ejfdZ!ejfdZ"xZ#S)RunLengthEncodingzi1D run length encoding. See `trimesh.voxel.runlength` documentation for implementation details. ct|tj|||jj }t |jjdk7r td||_ y)z Parameters ------------ data: run length encoded data. dtype: dtype of encoded data. Each second value of data is cast will be cast to this dtype if provided. rNrzdata must be 1D numpy array) rr%rrrrrr rz_dtype)r#r$rrs rr%zRunLengthEncoding.__init__s^ g33D9: =JJ$$E tzz A %:; ; rc tjtj|jddd|jddd SNrr)rrM logical_andrr)s rrFzRunLengthEncoding.is_emptys:66"..CaC$**QTT:JKLLLrcyrfr(r)s rrTzRunLengthEncoding.ndimssrc|jfSr'r/r)s rr zRunLengthEncoding.shape |rc|jSr')rr)s rrzRunLengthEncoding.dtyperrc6|jjSrarbr)s rr!zRunLengthEncoding.__hash__rcrcFttj|||SNry)rrr{)rrencoding_dtypes rrzRunLengthEncoding.from_denses!  " ":^ DE  rcd||jk7rtj||}t|Sr)rr rle_to_rlerrle_datars rfrom_rlezRunLengthEncoding.from_rles+ HNN " ++HEBH **rcBttj||Sr)rr brle_to_rle brle_datars r from_brlezRunLengthEncoding.from_brles !6!6y!NOOrc|jrt|jStj|j \}}|dk(r|}nt ||j}tj|d}||fS)NrrryrrH) rFrr r rle_striprrrrrr#r$rrs rrZzRunLengthEncoding.strippedsh =="4::. .!++DJJ7 g f H(T[[AH..q1  rch|jddd|jdddzjSr rr)s rr-zRunLengthEncoding.sums1 3Q3$**QTT"227799rc@tj|jSr')r rle_lengthrr)s rr/zRunLengthEncoding.sizes##DJJ//rcx|dk7rtd|ttj|jSNr z%encoding is 1D - cannot flip on axis )rzrr rle_reverserr]s rr_zRunLengthEncoding._flips6 4<DTHMN N !6!6tzz!BCCrc@tj|jSr')r rle_to_sparserr)s rrrz#RunLengthEncoding.sparse_components&&tzz22rc |jdSrErr)s rr1z RunLengthEncoding.sparse_indices rrc |jdSrfrr)s rr3zRunLengthEncoding.sparse_values rrcXtj|j|jSr)r rle_to_denserrr)s rr5zRunLengthEncoding.denses%%djj DDrcZtj|j||jSr)r rle_gather_1drrr7s rrzRunLengthEncoding.gathers&&tzz7$++NNrcRtj|d}|j|Srrrrr7s rr9zRunLengthEncoding.gather_nds!**W2.{{7##rctjttj|j ||j Sr)rrSrKrsorted_rle_gather_1drr)r#ordered_indicess r sorted_gatherzRunLengthEncoding.sorted_gathers3xx )00_M N++  rctjttj|j ||j Sr)rrSrKrrle_maskrrr<s rr=zRunLengthEncoding.mask"s-xxi00TBC4;;WWrcv|j|fD]$}tj||jcSyr)r7r asanyarrayr)r#r@rs rrAzRunLengthEncoding.get_value%s.''1E==dkk: :2rc`t|jj|jSr)rrrCrr)s rrCzRunLengthEncoding.copy)s !2$**EErcDtj|j|Sr)rrrr|s rr}z!RunLengthEncoding.run_length_data,s##DJJe<rrceZdZdZfdZej dZee jfdZ eddZ eddZ ej dZej dZej d Zd Zed Zej d Zej d Zej dZdZdZdZdZdZe jfdZe jfdZxZS)BinaryRunLengthEncodingzp1D binary run length encoding. See `trimesh.voxel.runlength` documentation for implementation details. c0t||ty)z_ Parameters ------------ data: binary run length encoded data. )r$rN)rr%rrs rr%z BinaryRunLengthEncoding.__init__9s d$/rcNtj|jddd SNrrrr)s rrFz BinaryRunLengthEncoding.is_emptyAs#66$**QTT*+++rcBttj||Sr)rBrr)rrs rrz"BinaryRunLengthEncoding.from_denseEs&  # #Jn E  rcBttj||Sr)rBrr?rs rrz BinaryRunLengthEncoding.from_rleKs&y'<'>rr')rrrrr%rrrFrrrrrrrZr-r/r_rrrr3r1r5rr9r7r=rCr}rrrs@rrBrB3sW 0 ,,.0hh  UU22   ! ! && 11K 77 66 44 33=$00:%'HH>,.88?rrBceZdZdZej dZej dZedZ edZ edZ edZ edZ ed Zd Zd Zy ) LazyIndexMapa8 Abstract class for implementing lazy index mapping operations. Implementations include transpose, flatten/reshaping and flipping Derived classes must implement: * _to_base_indices(indices) * _from_base_indices(base_indices) * shape * dense * mask(mask) cyr'r(r7s r_to_base_indiceszLazyIndexMap._to_base_indicesr:rcyr'r(r# base_indicess r_from_base_indiceszLazyIndexMap._from_base_indicesr:rc.|jjSr')rrFr)s rrFzLazyIndexMap.is_emptyszz"""rc.|jjSr'rr)s rrzLazyIndexMap.dtyperrc.|jjSr'rr)s rr-zLazyIndexMap.sumszz~~rc.|jjSr'rr)s rr/zLazyIndexMap.sizerrcL|j|jjSr')rhrr1r)s rr1zLazyIndexMap.sparse_indicess&&tzz'@'@AArc.|jjSr')rr3r)s rr3zLazyIndexMap.sparse_valuesszz'''rcV|jj|j|Sr')rr9rdr7s rr9zLazyIndexMap.gather_nds"zz##D$9$9'$BCCrcP|jt|j|Sr')rrKrdr?s rrAzLazyIndexMap.get_values!zz% 5 5e <=>>rN)rrrrrrrdrhrrFrr-r/r1r3r9rAr(rrrbrbs       ##  BB((D?rrbcXeZdZdZdZdZedZedZdZ edZ dZ y ) rlzr Lazily flattened encoding. Dense equivalent is np.reshape(data, (-1,)) (np.flatten creates a copy). c|tjtj||jjSr')rrrrr r7s rrdz"FlattenedEncoding._to_base_indicess(r//9I9IJKKrctjtj|j|jj dSr)rrrrrr rfs rrhz$FlattenedEncoding._from_base_indicess3~~  1A1A B  rc|jfSr'rr)s rr zFlattenedEncoding.shaperrcL|jjjdSr)rr5rjr)s rr5zFlattenedEncoding.denseszz''..rc~|jj|j|jjSr')rr=rjr r<s rr=zFlattenedEncoding.masks(zzt||DJJ,<,<=>>rc|Sr'r(r)s rrgzFlattenedEncoding.flats rcHt|jjSr')rlrrCr)s rrCzFlattenedEncoding.copys !233rN) rrrrrdrhrr r5r=rgrCr(rrrlrlsX L //?4rrlcjeZdZdZfdZdZdZedZedZ edZ dZ d Z xZ S) rhzI Lazily reshaped encoding. Numpy equivalent is `np.reshape` clt|tr|jdk7r|j}n t dt ||t||_|jjd}tj|j}|dk(rtj|}|jj|zdk7r4t d|jjt|j|jj|ztfd|jD|_y|dkDr t d tj|j|jjk7r4t d|jjt|jy) Nrzencoding must be an Encodingrrrz0cannot reshape encoding of size %d into shape %sc30K|] }|dk(rn|yw)rNr().0srems r z*ShapedEncoding.__init__.. sK{!qBwA 5{srz(shape cannot have more than one -1 value)rrrTrgrzrr%rKrcountrrabsrr/str)r#rr nnr/r~rs @rr%zShapedEncoding.__init__sE h )~~"#==;< < h'El [[  r "wwt{{# 766$.?sA&@19&@szperm z is not a valid permutation)rrrzrr rTrrr%rrSrrrL zeros_likearange_perm _inv_perm)r# base_encodingrinv_permrs ` rr%zTransposedEncoding.__init__3s-2F}FWXY Y }"" #s4y 0L##D   'xxBHH-AeM,?,?&@AAuTH,GHI I==&=#6#67 !rcnt|j|Dcgc]}|j|c}Scc}wr')rrr)r#rps rrzTransposedEncoding.transposeFs-4::t'Dt! 1 t'DEE'Ds2 ctdNzShould not be here RuntimeErrorrs rrzTransposedEncoding._transposeI/00rc|jSr')rr)s rrzTransposedEncoding.permLrvrcl|jjtfd|jDS)Nc3(K|] }| ywr'r()r|rr s rrz+TransposedEncoding.shape..Ss2z!U1Xzs)rr rKrris @rr zTransposedEncoding.shapePs'   2tzz222rcFtj||jdSr)rtakerr7s rrdz#TransposedEncoding._to_base_indicesUswww 44rc tj||jdS#t$r_tj|j tj |jj tj dcYSwxYwr)rrr TypeErrorastyperrfs rrhz%TransposedEncoding._from_base_indicesXsh 77<bA A 77##BHH-t~~/D/DRXX/NUW  s!$A%B  B c`|jjj|jSr')rr5rrr)s rr5zTransposedEncoding.denseas!zz))$**55rcV|jj|j|Sr')rr _base_indicesr7s rrzTransposedEncoding.gatheres"zz  !3!3G!<==rc|jj|j|jj|jSr')rr=rrrr<s rr=zTransposedEncoding.maskhs1zzt~~dnn=>HHTTrcP|jt|j|Sr')rrKrr?s rrAzTransposedEncoding.get_valueks!zz% 2 25 9:;;rc|jSr'rur)s rr$zTransposedEncoding.datanrvrc`t|jj|jS)N)rr)rrrCrr)s rrCzTransposedEncoding.copyrs! 0A SSr)rrrrr%rrrrr rdrhr5rr=rAr$rCrrs@rrr,s "&F133566>U<TrrcheZdZdZfdZdZdZedZedZ dZ dZ d d Z d Z xZS) r\zb Encoding with entries flipped along one or more axes. Dense equivalent is `np.flip` cJ|jt|tjr!|jdk(r|j f}nt|t r|f}tfd|D}tt|_ ttjtjk7rtdjt9|tfd|Ds*tdt!|j"jy)Nrc36K|]}|dkr|zn|ywrNr(r|rrTs rrz+FlippedEncoding.__init__..s!=1!a%QYQ.z$Axes cannot contain duplicates, got c3fK|](}d|cxkxrjjknc*ywr)rrT)r|rr#s rrz+FlippedEncoding.__init__..s);d1,DJJ,,,,ds.1z!Invalid axes %s for %d-d encoding)rTrrrr/itemrrKsorted_axesrsetrzrr%rrr)r#rr^rTrs` @rr%zFlippedEncoding.__init__}s dBJJ 'DIINIIK>D c "7D===6$<( s4:: 3tzz? 2CDJJ>RS S ";d;;3SY @P@P  11rr\c8t|ds|f}t}|jtfd|D}|D])}||vr|j ||j |+t |dk(r|S|jtt|S)N__iter__c36K|]}|dkr|zn|ywrr(rs rrz_flipped..!9Dqa!eU*Drr) hasattrrrTrKremoveaddrr_r)rr^ unique_axrrTs @rrnrns 4 $wI NNE 9D9 9D  >   Q  MM!    9~~~eF9$5677rc|jtfd|D}tjtj|k(r|S|j |S)Nc36K|]}|dkr|zn|ywrr()r|rrTs rrz_transposed..rr)rTrKrrrr)rrrTs @rrrsL NNE 9D9 9D vvbii$&'""4((rr')rrnumpyrrutilrrscipyrr BaseExceptionE exceptionsr rrrrrrrBrbrlrhrr\rnrr(rrrs= "]#s]#@M0HM0`MOXMO`Z"D>D>Na?/a?H2?82?j4 4D;L\;L|GTGTTC1lC1L8")A- ! BsB B=%B88B=