\ciE ddlZddlZddlZddlZddlmZddlmZddl m Z ddl m cm ZddlmZddlmZmZddlmZmZd gZd Zed d d gGdd ZdZejfdZdejddZdS)N)prod) GenericAlias)_dierckx) csr_array)array_namespacexp_capabilities) _not_a_knotBSpline NdBSplinecptj|tjr tjStjS)z>Return np.complex128 for complex dtypes, np.float64 otherwise.)np issubdtypecomplexfloating complex128float64dtypes l/var/www/html/mdtn/previsions/meteo_cartes/venv/lib/python3.11/site-packages/scipy/interpolate/_ndbspline.py _get_dtypers) }UB.//}zTF)z dask.arrayz7https://github.com/data-apis/array-api-extra/issues/488)cpu_onlyjax_jit skip_backendsceZdZdZeeZdddZedZ edZ edZ dddd Z edd Z dd ZdZdS)r aTensor product spline object. The value at point ``xp = (x1, x2, ..., xN)`` is evaluated as a linear combination of products of one-dimensional b-splines in each of the ``N`` dimensions:: c[i1, i2, ..., iN] * B(x1; i1, t1) * B(x2; i2, t2) * ... * B(xN; iN, tN) Here ``B(x; i, t)`` is the ``i``-th b-spline defined by the knot vector ``t`` evaluated at ``x``. Parameters ---------- t : tuple of 1D ndarrays knot vectors in directions 1, 2, ... N, ``len(t[i]) == n[i] + k + 1`` c : ndarray, shape (n1, n2, ..., nN, ...) b-spline coefficients k : int or length-d tuple of integers spline degrees. A single integer is interpreted as having this degree for all dimensions. extrapolate : bool, optional Whether to extrapolate out-of-bounds inputs, or return `nan`. Default is to extrapolate. Attributes ---------- t : tuple of ndarrays Knots vectors. c : ndarray Coefficients of the tensor-product spline. k : tuple of integers Degrees for each dimension. extrapolate : bool, optional Whether to extrapolate or return nans for out-of-bounds inputs. Defaults to true. Methods ------- __call__ derivative design_matrix See Also -------- BSpline : a one-dimensional B-spline object NdPPoly : an N-dimensional piecewise tensor product polynomial N extrapolatect||\|_|_\|_|_t |g|Rj|_|d}t||_ tj||_ |jj d}|j j |krtd|dt|D]}|j|}|j|}|j d|z dz } |j j || krzNdBSpline.t..sV  ;z&NdBSpline.__call__..s3"7"7"7&'#$rx'8"8"7"7"7r)r#r)rr'rzerosint64r%r*r+r/r-anyfloatreshaper0r(rkindviewravelstridesrevaluate_ndbspliner$r!r"r&) r1xirCrr*xi_shape was_complexccc1r _strides_c1num_c_troutrLs @r__call__zNdBSpline.__call__s*w}Q  *K;'' :4'222BBBbh///Bw!||rx{d22 -2--!$&kk---...26{{ O !Mb!M!MNNNZ% ( ( (8 ZZHRL ) )  !" % % B<4  KKKTKKLL Lgm(C/ W  $47<4//#B WWU^^ZZ$%/ 0 0hhjjj"7"7"7"7+-:"7"7"7>@hHHH 8B<)"!%!%!%!#!,!$!)!,!%!2   hhtw}%%kk(3B3-$'-*>>??}}S!!!rTcBtj|t}|jd}t ||kr#t dt |d|dt |\}\}tfdt|D}|ddd z} tj | dddtj ddd } tj |||| \} } } t| | | fS) a Construct the design matrix as a CSR format sparse array. Parameters ---------- xvals : ndarray, shape(npts, ndim) Data points. ``xvals[j, :]`` gives the ``j``-th data point as an ``ndim``-dimensional array. t : tuple of 1D ndarrays, length-ndim Knot vectors in directions 1, 2, ... ndim, k : int B-spline degree. extrapolate : bool, optional Whether to extrapolate out-of-bounds values of raise a `ValueError` Returns ------- design_matrix : a CSR array Each row of the design matrix corresponds to a value in `xvals` and contains values of b-spline basis elements which are non-zero at this value. rrGz*Data and knots are inconsistent: len(t) = z for ndim = rc3@K|]}||z dz VdSrNrI)r@r3r.len_ts rrAz*NdBSpline.design_matrix..s4AAa1Q4!+AAAAAArrNr)rr%rQr)r/r+r r;r,cumprodrOcopyr _coloc_ndr)clsxvalsr-r.rr*r"r#c_shapecscstridesdataindicesindptrrds ` @r design_matrixzNdBSpline.design_matrixsB0 5...{2 q66T>>SVV  (:!Q'?'?$<"e AAAAAU4[[AAAAAQRR[4 :b2hbh777"=BBDD!) 25E1lH!6!6gv$0111rrc tj||d}|jd}|jdd}||d}g} d} t |jdD]} ||krgt j||dd| f|} | |} | jdt| j | j z dz | _n8t j|tj t|dz d} | | j } | | jtj| dt| df|z}tj|d|}|| fS)NrrrG)axis)rmoveaxisr)rRr,r construct_fast derivativer2r/r-r.rNappendstack)r1r2r-r.rsrCr6trailing_shapec_flat new_c_listnew_tibdbnew_cs r_bspline_derivative_along_axisz(NdBSpline._bspline_derivative_along_axis sr K4 # # GAJ1b!! v|A'' $ $ABww*1fQQQTlA>>\\"%%t1SYY-112+ArxA /C/CQGG}   bd # # # #!,,,44 A   !N 244 E1d++e|rcXtj|tj}tj}|jdks|jd|kr(td|dtjdt|dkrtd|fdtj jdD}tj }j }t|D]T\}}|dkr ||||||| \}||<t#|||z d||<Ut%t'fd |D|t'|j S) a{ Construct a new NdBSpline representing the partial derivative. Parameters ---------- nu : array_like of shape (ndim,) Orders of the partial derivatives to compute along each dimension. Returns ------- NdBSpline A new NdBSpline representing the partial derivative of the original spline. rrrrErFrz-derivative orders must be positive, got nu = cHg|]}j|dj|fSr:)r#r$r?s rrMz(NdBSpline.derivative..Ds/NNNOT[^O+,NNNr)rCc3BK|]}|VdSr:)r&)r@r-r1s rrAz'NdBSpline.derivative..Qs/??At}}Q//??????rr)rr%rOr/r-r*r)r+rPr,r#listr.r(rg enumeratermaxr r;r&r) r1rCnu_arrr*t_newk_newc_newrsr6s ` rrvzNdBSpline.derivative)sBbh///46{{ ;!  v|A$66)r))df++)))** * vz?? QOOOPP PONNNeDGM!.os"3332HN2&&333rrz len(t) = z != len(k) = rrrzSpline degree in dimension z cannot be negative.zKnot vector in dimension z must be one-dimensional.zNeed at least z knots for degree z in dimension zKnots in dimension z# must be in a non-decreasing order.z.Need at least two internal knots in dimension z should not have nans or infs.c3 K|] }|dzV dSrcrI)r@r5s rrAz%_preprocess_inputs..s&%%R"q&%%%%%%rc6g|]}tj|SrI)rr%)r@r-s rrMz&_preprocess_inputs..s * * *qRZ]] * * *rc,g|]}t|SrIr/)r@tis rrMz&_preprocess_inputs..s % % %SWW % % %rN) isinstancer;r+r/ TypeErrorrr%rOr,r)r*diffrPuniqueisfiniteall unravel_indexarangerTrgemptyrrQfillnan) r.t_tplr*r3r4r5r6r)ror"rdr#s rr r Ws eU # # 1 %111   u::D A  DI 3333328DDDA 1vv~~ASZZAASVVAAABBB u::D 4[[00 Za ! ! qT HQK" q  66+1+++,, , 7a<<222233 3 rAv::8adQh88!#883488899 9 GBKK!O " " 87177788 8 ryBq1uH&& ' '! + +0+,00011 1{2""$$ 0/1///00 0 0 %%1%%% % %Erye55u==G:gRX6668==??L + *E * * *E u::D % %u % % %E 4U$E 2 2 2BGGBFOOO 4[[)) %a1ns58}}n  JuBH - - -E lRK ''sA AAc tj|jtjr0t ||j|fi|}t ||j|fi|}|d|zzS|jdkr|jddkrptj |}t|jdD]?}|||dd|ffi|\|dd|f<}|dkrtd|d|d|d@|S|||fi|\}}|dkrtd|d |d|S) Ny?rrrz solver = z returns info =z for column rz returns info = ) rrrr _iter_solverealimagr*r) empty_liker,r+) ar~solver solver_argsrrresjinfos rrrsg  }QWb0111aff<< <<1aff<< <<bg~v{{qwqzA~~mAqwqz"" R RA$fQ!!!Q$??;??OC1Itqyy !PF!P!P!P!PA!P!P!PQQQ F1a//;// T 199>>>D>>>?? ? rrc t}tdD} tn#t$r f|zYnwxYwtD]]\}}tt j|} | |kr+t d| d|d|d|dzd ^tfdt|D} t jd tj Dt } t | | } d d kr| |j} t!| d |t!| |d f}||}|t$jkr$t)jt,|}d|vrd|d<|| |fi|}||| |d z}t| |S)aConstruct an interpolating NdBspline. Parameters ---------- points : tuple of ndarrays of float, with shapes (m1,), ... (mN,) The points defining the regular grid in N dimensions. The points in each dimension (i.e. every element of the `points` tuple) must be strictly ascending or descending. values : ndarray of float, shape (m1, ..., mN, ...) The data on the regular grid in n dimensions. k : int, optional The spline degree. Must be odd. Default is cubic, k=3 solver : a `scipy.sparse.linalg` solver (iterative or direct), optional. An iterative solver from `scipy.sparse.linalg` or a direct one, `sparse.sparse.linalg.spsolve`. Used to solve the sparse linear system ``design_matrix @ coefficients = rhs`` for the coefficients. Default is `scipy.sparse.linalg.gcrotmk` solver_args : dict, optional Additional arguments for the solver. The call signature is ``solver(csr_array, rhs_vector, **solver_args)`` Returns ------- spl : NdBSpline object Notes ----- Boundary conditions are not-a-knot in all dimensions. c34K|]}t|VdSr:r)r@xs rrAzmake_ndbspl..s(,,SVV,,,,,,rz There are z points in dimension z , but order z requires at least rz points per dimension.c3K|]9}ttj|t|V:dS)rN)r rr%rQ)r@r3r.pointss rrAzmake_ndbspl..sX$$"*VAYe<<.s@@@r@@@rrrrNratolgư>)r/r;rrr atleast_1dr+r,r% itertoolsproductrQr rqeliminate_zerosr)rrRsslspsolve functoolspartialr)rvaluesr.rrr*rYr3pointnumptsr-rjmatrv_shape vals_shapevalscoefs` ` r make_ndbsplrs> v;;D,,V,,,,,H A  DIf%%AA5R]5))** QqT>>@&@@q@@+,Q4@@!"1a@@@AA A  $$$$$T{{$$$ $ $A J@@Y%6%?@@@ N N NE  " "5!Q / /D tqyy  lGwuu~&&WTUU^(<(<=J >>* % %D ";v>>>  $ $"&K  6$ , , , ,D <<7455>1 2 2D Qa  s<AA)r)rrrnumpyrmathrtypesrrscipy.sparse.linalgsparselinalgr scipy.sparserscipy._lib._array_apirr _bsplinesr r __all__rr r gcrotmkrrrIrrrs!!!!!!!!!""""""BBBBBBBB++++++++ - 5 Dr r r r r r r r h J(J(J(Z![0J!s{J!J!J!J!J!J!J!r