Package: jitfields
jitfieldsFast functions for dense scalar and vector fields, implemented using just-in-time compilation./!\ This is (still) experimentalGPU version of the algorithms are written in pure CUDA, and compiled just-in-time bycupy.CPU version of the algorithms are written in pure C++, and compiled just-in-time bycppyy.InstallationInstallation through pip should work, although I don't know how robust the cupy/pytorch
interaction is in term of cuda version.pipinstallgit+https://github.com/balbasty/jitfieldsIf you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.pipinstallgit+https://github.com/balbasty/jitfields#egg=jitfields[cuda]Pre-installing dependencies from conda may be more robust:condainstall-cpython>=3.6conda-forgenumpycupyccpyypytorch>=1.8cudatoolkit=10.2
pipinstallgit+https://github.com/balbasty/jitfields#egg=jitfields[cuda]Implemented so farDistance transformsDistance to binary masksdefeuclidean_distance_transform(x,ndim=None,vx=1,dtype=None):..."""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensorInput tensorndim : int, default=`x.ndim`Number of spatial dimensionsvx : [sequence of] float, default=1Voxel sizeReturns-------d : (..., *spatial) tensorDistance mapReferences----------..[1] "Distance Transforms of Sampled Functions"Pedro F. Felzenszwalb & Daniel P. HuttenlocherTheory of Computing (2012)https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""defl1_distance_transform(x,ndim=None,vx=1,dtype=None):..."""Compute the L1 distance transform of a binary imageParameters----------x : (..., *spatial) tensorInput tensordim : int, default=`x.ndim`Number of spatial dimensionsvx : [sequence of] float, default=1Voxel sizedtype : torch.dtypeDatatype of the distance map.By default, use x.dtype if it is a floating point type,otherwise use the default floating point type.Returns-------d : (..., *spatial) tensorDistance mapReferences----------..[1] "Distance Transforms of Sampled Functions"Pedro F. Felzenszwalb & Daniel P. HuttenlocherTheory of Computing (2012)https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""defsigned_distance_transform(x:tensor,ndim:Optional[int]=None,vx:OneOrSeveral[float]=1,dtype:Optional[torch.dtype]=None,)->tensor:..."""Compute the signed Euclidean distance transform of a binary imageParameters----------x : `(..., *spatial) tensor`Input tensor, with shape `(..., *spatial)`.ndim : `int`, default=`x.ndim`Number of spatial dimensions. Default: all.vx : `[sequence of] float`, default=1Voxel size.dtype : `torch.dtype`, optionalOuptut data type. Default is same as `x` if it has a floatingpoint data type, else `torch.get_default_dtype()`.Returns-------d : `(..., *spatial) tensor`Signed distance map, with shape `(..., *spatial)`.References----------..[1] "Distance Transforms of Sampled Functions"Pedro F. Felzenszwalb & Daniel P. HuttenlocherTheory of Computing (2012)https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""Distance to 1D splinesdefspline_distance_table(loc:tensor,coeff:tensor,steps:Optional[Union[int,tensor]]=None,order:OrderType=3,bound:BoundType='dct2',square:bool=False,)->Tuple[tensor,tensor]:..."""Compute the minimum distance from a set of points to a 1D splineParameters----------loc : `(..., D) tensor`Point set.coeff : `(..., N, D) tensor`Spline coefficients encoding the location of the 1D spline.steps : `int or (..., K) tensor`Number of time steps to try, or list of time steps to try.order : {1..7}Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}`Boundary conditions of the spline.square : boolReturn the squared Euclidean distance.Returns-------dist : `(...) tensor`Distance from each point in the set to its closest point on the splinetime : `(...) tensor`Time of the closest point on the spline"""defspline_distance_brent(loc:tensor,coeff:tensor,max_iter:int=128,tol:float=1e-6,step_size:float=0.01,order:OrderType=3,bound:BoundType='dct2',square:bool=False,steps:Optional[Union[int,tensor]]=None,)->Tuple[tensor,tensor]:..."""Compute the minimum distance from a set of points to a 1D splineParameters----------loc : `(..., D) tensor`Point set.coeff : `(..., N, D) tensor`Spline coefficients encoding the location of the 1D spline.max_iter : intNumber of optimization steps.tol : floatTolerance for early stoppingstep_size : floatInitial search size.order : {1..7}Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}`Boundary conditions of the spline.square : boolReturn the squared Euclidean distance.steps : intNumber of steps used in the table-based initialisation.Returns-------dist : `(...) tensor`Distance from each point in the set to its closest point on the splinetime : `(...) tensor`Time of the closest point on the spline"""defspline_distance_gaussnewton(loc:tensor,coeff:tensor,max_iter:int=16,tol:float=1e-6,order:OrderType=3,bound:BoundType='dct2',square:bool=False,steps:Optional[Union[int,tensor]]=None,)->Tuple[tensor,tensor]:..."""Compute the minimum distance from a set of points to a 1D splineParameters----------loc : `(..., D) tensor`Point set.coeff : `(..., N, D) tensor`Spline coefficients encoding the location of the 1D spline.max_iter : intNumber of optimization steps.tol : floatTolerance for early stoppingorder : {1..7}Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}`Boundary conditions of the spline.square : boolReturn the squared Euclidean distance.steps : intNumber of steps used in the table-based initialisation.Returns-------dist : `(...) tensor`Distance from each point in the set to its closest point on the splinetime : `(...) tensor`Time of the closest point on the spline"""defspline_distance_brent_(dist:tensor,time:tensor,loc:tensor,coeff:tensor,max_iter:int=128,tol:float=1e-6,step_size:float=0.01,order:OrderType=3,bound:BoundType='dct2',square:bool=False,)->Tuple[tensor,tensor]:..."""Compute the minimum distance from a set of points to a 1D spline (inplace)Parameters----------dist : `(...) tensor`Initial distance from each point in the set to its closest point on the splinetime : `(...) tensor`Initial time of the closest point on the splineloc : `(..., D) tensor`Point set.coeff : `(..., N, D) tensor`Spline coefficients encoding the location of the 1D spline.max_iter : intNumber of optimization steps.tol : floatTolerance for early stoppingstep_size : floatInitial search size.order : {1..7}Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}`Boundary conditions of the spline.square : boolReturn the squared Euclidean distance.Returns-------dist : `(...) tensor`Distance from each point in the set to its closest point on the splinetime : `(...) tensor`Time of the closest point on the spline"""defspline_distance_gaussnewton_(dist:tensor,time:tensor,loc:tensor,coeff:tensor,max_iter:int=16,tol:float=1e-6,order:OrderType=3,bound:BoundType='dct2',square:bool=False,)->Tuple[tensor,tensor]:..."""Compute the minimum distance from a set of points to a 1D spline (inplace)Parameters----------dist : `(...) tensor`Initial distance from each point in the set to its closest point on the splinetime : `(...) tensor`Initial time of the closest point on the splineloc : `(..., D) tensor`Point set.coeff : `(..., N, D) tensor`Spline coefficients encoding the location of the 1D spline.max_iter : intNumber of optimization steps.tol : floatTolerance for early stoppingorder : {1..7}Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}`Boundary conditions of the spline.square : boolReturn the squared Euclidean distance.Returns-------dist : `(...) tensor`Distance from each point in the set to its closest point on the splinetime : `(...) tensor`Time of the closest point on the spline"""Distance to triangular meshesdefmesh_distance_signed(loc:tensor,vertices:tensor,faces:tensor,out:Optional[tensor]=None,)->tensor:..."""Compute the *signed* minimum distance from a set of points to a triangular meshParameters----------loc : `(..., D) tensor`Point set.vertices : `(N, D) tensor`Mesh verticesfaces : `(M, D) tensor[integer]`Mesh facesReturns-------dist : `(...) tensor`Signed distance from each point in the set to its closest point on the mesh(negative inside, positive outside)"""defmesh_distance(loc:tensor,vertices:tensor,faces:tensor,out:Optional[tensor]=None,)->tensor:..."""Compute the minimum distance from a set of points to a triangular meshParameters----------loc : `(..., D) tensor`Point set.vertices : `(N, D) tensor`Mesh verticesfaces : `(M, D) tensor[integer]`Mesh facesReturns-------dist : `(...) tensor`Signed distance from each point in the set to its closest point on the mesh(negative inside, positive outside)"""Interpolation/Resamplingdefspline_coeff(inp,order,bound='dct2',dim=-1):..."""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensorInput tensororder : {0..7}, default=2Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2'Boundary conditions.dim : int, default=-1Dimension along which to filterReturns-------coeff : tensorSpline coefficientsReferences----------..[1]  M. Unser, A. Aldroubi and M. Eden."B-Spline Signal Processing: Part I-Theory,"IEEE Transactions on Signal Processing 41(2):821-832 (1993)...[2]  M. Unser, A. Aldroubi and M. Eden."B-Spline Signal Processing: Part II-Efficient Design and Applications,"IEEE Transactions on Signal Processing 41(2):834-848 (1993)...[3]  M. Unser."Splines: A Perfect Fit for Signal and Image Processing,"IEEE Signal Processing Magazine 16(6):22-38 (1999)."""defspline_coeff_nd(inp,order,bound='dct2',ndim=None):..."""Compute the interpolating spline coefficients, along the last N dimensions.Parameters----------inp : (..., *spatial) tensorInput tensororder : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2'Boundary conditions.ndim : int, default=`inp.dim()`Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensorSpline coefficientsReferences----------..[1]  M. Unser, A. Aldroubi and M. Eden."B-Spline Signal Processing: Part I-Theory,"IEEE Transactions on Signal Processing 41(2):821-832 (1993)...[2]  M. Unser, A. Aldroubi and M. Eden."B-Spline Signal Processing: Part II-Efficient Design and Applications,"IEEE Transactions on Signal Processing 41(2):834-848 (1993)...[3]  M. Unser."Splines: A Perfect Fit for Signal and Image Processing,"IEEE Signal Processing Magazine 16(6):22-38 (1999)."""defresize(x,factor=None,shape=None,ndim=None,anchor='e',order=2,bound='dct2',prefilter=True):..."""Resize a tensor using spline interpolationParameters----------x : (..., *inshape) tensorInput  tensorfactor : [sequence of] float, optionalFactor by which to resize the tensor (> 1 == bigger)One of factor or shape must be provided.shape : [sequence of] float, optionalShape of output tensor.One of factor or shape must be provided.ndim : int, optionalNumber if spatial dimensions.If not provided, try to guess from factor or shape.If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or NoneWhat feature should be aligned across the input and output tensors.If 'edge' or 'center', the effective scaling factor may slightlydiffer from the requested scaling factor.If None, the center of the (0, 0) voxel is aligned, and therequested factor is exactly applied.order : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2'How to deal with out-of-bound values.prefilter : bool, default=TrueWhether to first compute interpolating coefficients.Must be true for proper interpolation, otherwise thisfunction merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensorResized tensorReferences----------..[1]  M. Unser, A. Aldroubi and M. Eden."B-Spline Signal Processing: Part I-Theory,"IEEE Transactions on Signal Processing 41(2):821-832 (1993)...[2]  M. Unser, A. Aldroubi and M. Eden."B-Spline Signal Processing: Part II-Efficient Design and Applications,"IEEE Transactions on Signal Processing 41(2):834-848 (1993)...[3]  M. Unser."Splines: A Perfect Fit for Signal and Image Processing,"IEEE Signal Processing Magazine 16(6):22-38 (1999)."""defrestrict(x,factor=None,shape=None,ndim=None,anchor='e',order=2,bound='dct2',reduce_sum=False):..."""Restrict (adjoint of resize) a tensor using spline interpolationParameters----------x : (..., *inshape) tensorInput  tensorfactor : [sequence of] float, optionalFactor by which to resize the tensor (> 1 == smaller)One of factor or shape must be provided.shape : [sequence of] float, optionalShape of output tensor.One of factor or shape must be provided.ndim : int, optionalNumber if spatial dimensions.If not provided, try to guess from factor or shape.If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or NoneWhat feature should be aligned across the input and output tensors.If 'edge' or 'center', the effective scaling factor may slightlydiffer from the requested scaling factor.If None, the center of the (0, 0) voxel is aligned, and therequested factor is exactly applied.order : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2'How to deal with out-of-bound values.Returns-------x : (..., *shape) tensorrestricted tensor"""defpull(inp,grid,order=2,bound='dct2',extrapolate=True,prefilter=False,out=None):..."""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensorInput tensorgrid : (..., *outshape, ndim) tensorTensor of coordinates into `inp`order : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2'How to deal with out-of-bound values.extrapolate : bool or {'center', 'edge'}- True: use bound to extrapolate out-of-bound value- False or 'center': do not extrapolate values that fall outsideof the centers of the first and last voxels.- 'edge': do not extrapolate values that fall outsideof the edges of the first and last voxels.prefilter : bool, default=TrueWhether to first compute interpolating coefficients.Must be true for proper interpolation, otherwise thisfunction merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensorPulled tensor"""defpush(inp,grid,shape=None,order=2,bound='dct2',extrapolate=True,prefilter=False,out=None):..."""Splat a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensorInput tensorgrid : (..., *inshape, ndim) tensorTensor of coordinates into `inp`shape : sequence[int], default=inshapeOutput spatial shapeorder : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2'How to deal with out-of-bound values.extrapolate : bool or {'center', 'edge'}- True: use bound to extrapolate out-of-bound value- False or 'center': do not extrapolate values that fall outsideof the centers of the first and last voxels.- 'edge': do not extrapolate values that fall outsideof the edges of the first and last voxels.prefilter : bool, default=TrueWhether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensorPulled tensor"""defcount(grid,shape=None,order=2,bound='dct2',extrapolate=True,out=None):..."""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensorTensor of coordinatesshape : sequence[int], default=inshapeOutput spatial shapeorder : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2'How to deal with out-of-bound values.extrapolate : bool or {'center', 'edge'}- True: use bound to extrapolate out-of-bound value- False or 'center': do not extrapolate values that fall outsideof the centers of the first and last voxels.- 'edge': do not extrapolate values that fall outsideof the edges of the first and last voxels.Returns-------out : (..., *shape) tensorPulled tensor"""defgrad(inp,grid,order=2,bound='dct2',extrapolate=True,prefilter=False,out=None):..."""Sample the spatial gradients of a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensorInput tensorgrid : (..., *outshape, ndim) tensorTensor of coordinates into `inp`order : [sequence of] {0..7}, default=2Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2'How to deal with out-of-bound values.extrapolate : bool or {'center', 'edge'}- True: use bound to extrapolate out-of-bound value- False or 'center': do not extrapolate values that fall outsideof the centers of the first and last voxels.- 'edge': do not extrapolate values that fall outsideof the edges of the first and last voxels.prefilter : bool, default=TrueWhether to first compute interpolating coefficients.Must be true for proper interpolation, otherwise thisfunction merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensorPulled gradients"""Compact symmetric (or postive-definite) matricesdefsym_matvec(mat,vec,dtype=None,out=None):..."""Matrix-vector product for compact symmetric matrices`out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVectordtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.out : (..., C) tensor, optionalOutput placeholderReturns-------out : (..., C) tensorMatrix-vector product"""defsym_addmatvec(inp,mat,vec,dtype=None,out=None):..."""Add a matrix-vector product for compact symmetric matrices`out = inp + mat @ vec`Parameters----------inp : (..., C) tensorVector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVector used in the matrix-vector productdtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.out : (..., C) tensor, optionalOutput placeholderReturns-------out : (..., C) tensorAdded matrix-vector product"""defsym_addmatvec_(inp,mat,vec,dtype=None):..."""Inplace add a matrix-vector product for compact symmetric matrices`inp += mat @ vec`Parameters----------inp : (..., C) tensorVector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVector used in the matrix-vector productdtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensorAdded matrix-vector product"""defsym_submatvec(inp,mat,vec,dtype=None,out=None):..."""Subtract a matrix-vector product for compact symmetric matrices`out = inp - mat @ vec`Parameters----------inp : (..., C) tensorVector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVector used in the matrix-vector productdtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.out : (..., C) tensor, optionalOutput placeholderReturns-------out : (..., C) tensorSubtracted matrix-vector product"""defsym_submatvec_(inp,mat,vec,dtype=None):..."""Inplace subtract a matrix-vector product for compact symmetric matrices`inp -= mat @ vec`Parameters----------inp : (..., C) tensorVector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVector used in the matrix-vector productdtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensorSubtracted matrix-vector product"""defsym_solve(mat,vec,dtype=None,out=None):..."""Solve the symmetric linear system`out = mat.inverse() @ vec`!! Does not backpropagate through `mat` !!Parameters----------mat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVectordtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.out : (..., C) tensor, optionalOutput placeholderReturns-------out : (..., C) tensorSolution of the linear system"""defsym_solve_(mat,vec,dtype=None):..."""Solve the symmetric linear system in-place`vec = mat.inverse() @ vec`!! Does not backpropagate through `mat` !!Parameters----------mat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.vec : (..., C) tensorVectordtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensorSolution of the linear system"""defsym_invert(mat,dtype=None,out=None):..."""Invert a compact symmetric matrix`out = mat.inverse()`!! Does not backpropagate through `mat` !!Parameters----------mat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.dtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optionalOutput placeholderReturns-------mat : (..., C*(C+1)//2) tensorInverse matrix"""defsym_invert_(mat,dtype=None):..."""Invert a compact symmetric matrix in-place`mat = mat.inverse()`!! Does not backpropagate through `mat` !!Parameters----------mat : (..., C*(C+1)//2) tensorSymmetric matrix with compact storage.The matrix should be saved as a vector containing the diagonalfollowed by the rows of the upper triangle.dtype : torch.dtype, optionalData type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensorInverse matrix"""Regularisers for dense flow fieldsdefflow_matvec(vel:Tensor,weight:Optional[Tensor]=None,absolute:float=0,membrane:float=0,bending:float=0,shears:float=0,div:float=0,bound:list[str]='dft',voxel_size:list[float]=1,out:Optional[Tensor]=None)->Tensor:..."""Apply a spatial regularization matrix.Parameters----------vel : (*batch, *spatial, ndim) tensorInput displacement field, in voxels.weight : (*batch, *spatial) tensor, optionalWeight map, to spatially modulate the regularization.absolute : floatPenalty on absolute values.membrane : floatPenalty on first derivatives.bending : floatPenalty on second derivatives.shears : floatPenalty on local shears.div : floatPenalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft'Boundary conditions.voxel_size : [sequence of] floatVoxel size.out : (*batch, *spatial, ndim) tensor, optionalOutput placeholderReturns-------out : (*batch, *spatial, ndim) tensor"""# We also implement variants that adds to or subtracts from an input tensordefflow_matvec_add(inp:Tensor,...):...defflow_matvec_add_(inp:Tensor,...):...defflow_matvec_sub(inp:Tensor,...):...defflow_matvec_sub_(inp:Tensor,...):...defflow_kernel(shape:list[int],absolute:float=0,membrane:float=0,bending:float=0,shears:float=0,div:float=0,bound:list[str]='dft',voxel_size:list[float]=1,out:Optional[Tensor]=None)->Tensor:..."""Return the kernel of a Toeplitz regularization matrix.Parameters----------shape : int or list[int]Number of spatial dimensions or shape of the tensorabsolute : floatPenalty on absolute values.membrane : floatPenalty on first derivatives.bending : floatPenalty on second derivatives.shears : floatPenalty on local shears. Linear elastic energy's `mu`.div : floatPenalty on local volume changes. Linear elastic energy's `lambda`.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft'Boundary conditions.voxel_size : [sequence of] floatVoxel size.out : (*shape, ndim, [ndim]) tensor, optionalOutput placeholderReturns-------out : (*shape, ndim, [ndim]) tensorConvolution kernel.A matrix or kernels ([ndim, ndim]) if `shears` or `div`,else a vector of kernels ([ndim]) ."""# We also implement variants that adds to or subtracts from an input tensordefflow_kernel_add(inp:Tensor,...):...defflow_kernel_add_(inp:Tensor,...):...defflow_kernel_sub(inp:Tensor,...):...defflow_kernel_sub_(inp:Tensor,...):...defflow_diag(shape:list[int],weight:Optional[Tensor]=None,absolute:float=0,membrane:float=0,bending:float=0,shears:float=0,div:float=0,bound:list[str]='dft',voxel_size:list[float]=1,out:Optional[Tensor]=None)->Tensor:..."""Return the diagonal of a regularization matrix.Parameters----------shape : list[int]Shape of the tensorweight : (*batch, *spatial) tensor, optionalWeight map, to spatially modulate the regularization.absolute : floatPenalty on absolute values.membrane : floatPenalty on first derivatives.bending : floatPenalty on second derivatives.shears : floatPenalty on local shears.div : floatPenalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft'Boundary conditions.voxel_size : [sequence of] floatVoxel size.out : (*batch, *spatial, ndim) tensor, optionalOutput placeholderReturns-------out : (*batch, *spatial, ndim) tensor"""# We also implement variants that adds to or subtracts from an input tensordefflow_diag_add(inp:Tensor,...):...defflow_diag_add_(inp:Tensor,...):...defflow_diag_sub(inp:Tensor,...):...defflow_diag_sub_(inp:Tensor,...):...defflow_relax_(vel:Tensor,hes:Tensor,grd:Tensor,weight:Optional[Tensor]=None,absolute:float=0,membrane:float=0,bending:float=0,shears:float=0,div:float=0,bound:list[str]='dft',voxel_size:list[float]=1,nb_iter:int=1,)->Tensor:..."""Perform relaxation iterations.Parameters----------vel : (*batch, *spatial, ndim) tensorWarm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensorInput symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensorInput gradient, in voxels.weight : (*batch, *spatial) tensor, optionalWeight map, to spatially modulate the regularization.absolute : floatPenalty on absolute values.membrane : floatPenalty on first derivatives.bending : floatPenalty on second derivatives.shears : floatPenalty on local shears.div : floatPenalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft'Boundary conditions.voxel_size : [sequence of] floatVoxel size.nb_iter : intNumber of iterationsReturns-------vel : (*batch, *spatial, ndim) tensor"""defflow_precond(mat:Tensor,vec:Tensor,weight:Optional[Tensor]=None,absolute:float=0,membrane:float=0,bending:float=0,shears:float=0,div:float=0,bound:list[str]='dft',voxel_size:list[float]=1,out:Optional[Tensor]=None)->Tensor:..."""Apply the preconditioning `(M + diag(R)) \ v`Parameters----------mat : (*batch, *spatial, DD) tensorDD == 1 | D | D*(D+1)//2 | D*DPreconditioning matrix `M`vec : (*batch, *spatial, D) tensorPoint `v` at which to solve the system.weight : (*batch, *spatial) tensor, optionalRegularization weight map.absolute : floatPenalty on absolute values.membrane : floatPenalty on first derivatives.bending : floatPenalty on second derivatives.shears : floatPenalty on local shears.div : floatPenalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft'Boundary conditions.voxel_size : [sequence of] floatVoxel size.out : (*batch, *spatial, D) tensorOutput placeholder.Returns-------out : (*batch, *spatial, D) tensorPreconditioned vector."""defflow_forward(mat:Tensor,vec:Tensor,weight:Optional[Tensor]=None,absolute:float=0,membrane:float=0,bending:float=0,shears:float=0,div:float=0,bound:list[str]='dft',voxel_size:list[float]=1,out:Optional[Tensor]=None)->Tensor:..."""Apply the forward matrix-vector product `(M + R) @ v`Parameters----------mat : (*batch, *spatial, DD) tensorDD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensorPoint `v` at which to solve the system.weight : (*batch, *spatial) tensor, optionalRegularization weight map.absolute : floatPenalty on absolute values.membrane : floatPenalty on first derivatives.bending : floatPenalty on second derivatives.shears : floatPenalty on local shears.div : floatPenalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft'Boundary conditions.voxel_size : [sequence of] floatVoxel size.out : (*batch, *spatial, D) tensorOutput placeholder.Returns-------out : (*batch, *spatial, D) tensorPreconditioned vector."""