BranchesTags

Repository files navigation

jitfields

Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.

/!\ This is (still) experimental

  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.

Installation

Dependencies

  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)

Conda

PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves in situation of incompatibility. The preferred installation method therefore relies on conda, which minimizes such issues.

conda install jitfields -c balbasty -c pytorch -c conda-forge

Note that in this case PyTorch without GPU support will get installed (unless PyTorch was already installed using conda, in which case the installed version will be preserved). To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy), you should instead do:

# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge

In our experience this is enough to ensure compatibility across all dependencies. If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:

Pip

Installation through pip should work, as jitfields is a pure python package. As stated above, there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable to pre-install these dependencies yourself, rather than relying on pip's dependency solver.

pip install jitfields

If you intend to run code on the GPU, specify the [cuda] extra tag, which ensures that cupy gets installed.

pip install jitfields[cuda]

API

Distance transforms

Distance to binary masks

defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""

Distance to 1D splines

defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes

defmesh_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/Resampling

defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""

Compact symmetric (or postive-definite) matrices

defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""

Regularisers for dense flow fields

defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""

About

Fast functions for dense scalar and vector fields, implemented using just-in-time compilation

jitfields.readthedocs.io

Resources

Readme
MIT license
Activity

Stars

2 stars

Watchers

2 watching

Forks

2 forks
Report repository

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
BranchesTags
Repository files navigation
jitfields
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation.
/!\ This is (still) experimental
  • GPU version of the algorithms are written in pure CUDA, and compiled just-in-time by cupy.
  • CPU version of the algorithms are written in pure C++, and compiled just-in-time by cppyy.
Installation
Dependencies
  • pytorch >= 1.8
  • numpy
  • cppyy
  • cupy (if CUDA support required)
Conda
PyTorch, cppyy and cupy all heavily depend on system libraries, and easily find themselves
in situation of incompatibility. The preferred installation method therefore relies on
conda, which minimizes such issues.
conda install jitfields -c balbasty -c pytorch -c conda-forge 
Note that in this case PyTorch without GPU support will get installed (unless PyTorch
was already installed using conda, in which case the installed version will be preserved).
To ensure that the GPU version of PyTorch gets installed (and ensure compatibility with cupy),
you should instead do:
# for pytorch >= 1.13
conda install jitfields pytorch==$TORCH_VERSION pytorch-cuda=$CUDA_VERSION -c balbasty -c pytorch -c nvidia -c conda-forge # for pytorch < 1.13
conda install jitfields pytorch==$TORCH_VERSION cudatoolkit=$CUDA_VERSION -c balbasty -c pytorch -c conda-forge 
In our experience this is enough to ensure compatibility across all dependencies.
If for some reason it is not, it may be necessary to use cupy's specific cuda-version package. See:
Pip
Installation through pip should work, as jitfields is a pure python package. As stated above,
there may be inconsistencies across pytorch, cppyy and cupy. It may therefore be preferable
to pre-install these dependencies yourself, rather than relying on pip's dependency solver.
pip install jitfields
If you intend to run code on the GPU, specify the [cuda] extra tag, which
ensures that cupy gets installed.
pip install jitfields[cuda]
API
Distance transforms
Distance to binary masks
defeuclidean_distance_transform(x, ndim=None, vx=1, dtype=None): ...
"""Compute the Euclidean distance transform of a binary imageParameters----------x : (..., *spatial) tensor Input tensorndim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizeReturns-------d : (..., *spatial) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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) tensor Input tensordim : int, default=`x.ndim` Number of spatial dimensionsvx : [sequence of] float, default=1 Voxel sizedtype : torch.dtype Datatype 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) tensor Distance mapReferences----------..[1] "Distance Transforms of Sampled Functions" Pedro F. Felzenszwalb & Daniel P. Huttenlocher Theory 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=1 Voxel size.dtype : `torch.dtype`, optional Ouptut data type. Default is same as `x` if it has a floating point 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. Huttenlocher Theory of Computing (2012) https://www.theoryofcomputing.org/articles/v008a019/v008a019.pdf"""
Distance to 1D splines
defspline_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 : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return the squared Euclidean distance.steps : int Number 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 : int Number of optimization steps.tol : float Tolerance for early stoppingstep_size : float Initial search size.order : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 : int Number of optimization steps.tol : float Tolerance for early stoppingorder : {1..7} Spline order.bound : `{'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}` Boundary conditions of the spline.square : bool Return 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 meshes
defmesh_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/Resampling
defspline_coeff(inp, order, bound='dct2', dim=-1): ...
"""Compute the interpolating spline coefficients, along a single dimension.Parameters----------inp : tensor Input tensororder : {0..7}, default=2 Interpolation order.bound : {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.dim : int, default=-1 Dimension along which to filterReturns-------coeff : tensor Spline 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) tensor Input tensororder : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dft'}, default='dct2' Boundary conditions.ndim : int, default=`inp.dim()` Number of spatial dimensionsReturns-------coeff : (..., *spatial) tensor Spline 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == bigger) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "prolongation".Returns-------x : (..., *shape) tensor Resized 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) tensor Input tensorfactor : [sequence of] float, optional Factor by which to resize the tensor (> 1 == smaller) One of factor or shape must be provided.shape : [sequence of] float, optional Shape of output tensor. One of factor or shape must be provided.ndim : int, optional Number if spatial dimensions. If not provided, try to guess from factor or shape. If guess fails, assume ndim = x.dim().anchor : {'edge', 'center'} or None What feature should be aligned across the input and output tensors. If 'edge' or 'center', the effective scaling factor may slightly differ from the requested scaling factor. If None, the center of the (0, 0) voxel is aligned, and the requested factor is exactly applied.order : [sequence of] {0..7}, default=2 Interpolation order.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dct2' How to deal with out-of-bound values.Returns-------x : (..., *shape) tensor restricted tensor"""
defpull(inp, grid, order=2, bound='dct2', extrapolate=True, prefilter=False, out=None): ...
"""Sample a tensor using spline interpolationParameters----------inp : (..., *inshape, channel) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel) tensor Pulled 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) tensor Input tensorgrid : (..., *inshape, ndim) tensor Tensor of coordinates into `inp`shape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels. prefilter : bool, default=True Whether to compute interpolating coefficients at the end.Returns-------out : (..., *shape, channel) tensor Pulled tensor"""
defcount(grid, shape=None, order=2, bound='dct2', extrapolate=True, out=None): ...
"""Splat ones using spline interpolationParameters----------grid : (..., *inshape, ndim) tensor Tensor of coordinatesshape : sequence[int], default=inshape Output spatial shapeorder : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.Returns-------out : (..., *shape) tensor Pulled 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) tensor Input tensorgrid : (..., *outshape, ndim) tensor Tensor of coordinates into `inp`order : [sequence of] {0..7}, default=2 Interpolation 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 outside of the centers of the first and last voxels. - 'edge': do not extrapolate values that fall outside of the edges of the first and last voxels.prefilter : bool, default=True Whether to first compute interpolating coefficients. Must be true for proper interpolation, otherwise this function merely performs a non-interpolating "spline sampling".Returns-------out : (..., *outshape, channel, ndim) tensor Pulled gradients"""
Compact symmetric (or postive-definite) matrices
defsym_matvec(mat, vec, dtype=None, out=None): ...
"""Matrix-vector product for compact symmetric matrices `out = mat @ vec`Parameters----------mat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Matrix-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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Added 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Subtracted 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) tensor Vector to which the matrix-vector product is addedmat : (..., C*(C+1)//2) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vector used in the matrix-vector productdtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------inp : (..., C) tensor Subtracted 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C) tensor, optional Output placeholderReturns-------out : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.vec : (..., C) tensor Vectordtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------vec : (..., C) tensor Solution 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.out : (..., C*(C+1)//2) tensor, optional Output placeholderReturns-------mat : (..., C*(C+1)//2) tensor Inverse 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) tensor Symmetric matrix with compact storage. The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.dtype : torch.dtype, optional Data type used to carry the computation. By default, same as input.Returns-------mat : (..., C*(C+1)//2) tensor Inverse matrix"""
Regularisers for dense flow fields
defflow_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) tensor Input displacement field, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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 : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears. Linear elastic energy's `mu`.div : float Penalty 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] float Voxel size.out : (*shape, ndim, [ndim]) tensor, optional Output placeholderReturns-------out : (*shape, ndim, [ndim]) tensor Convolution 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, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, ndim) tensor, optional Output 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) tensor Warm start.hes : (*batch, *spatial, ndim*(ndim+1)//2) tensor Input symmetric Hessian, in voxels.grd : (*batch, *spatial, ndim) tensor Input gradient, in voxels.weight : (*batch, *spatial) tensor, optional Weight map, to spatially modulate the regularization.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.nb_iter : int Number 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) tensor DD == 1 | D | D*(D+1)//2 | D*D Preconditioning matrix `M`vec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned 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) tensor DD == 1 | D | D*(D+1)//2 | D*Dvec : (*batch, *spatial, D) tensor Point `v` at which to solve the system.weight : (*batch, *spatial) tensor, optional Regularization weight map.absolute : float Penalty on absolute values.membrane : float Penalty on first derivatives.bending : float Penalty on second derivatives.shears : float Penalty on local shears.div : float Penalty on local volume changes.bound : [sequence of] {'zero', 'replicate', 'dct1', 'dct2', 'dst1', 'dst2', 'dft'}, default='dft' Boundary conditions.voxel_size : [sequence of] float Voxel size.out : (*batch, *spatial, D) tensor Output placeholder.Returns-------out : (*batch, *spatial, D) tensor Preconditioned vector."""
About
Fast functions for dense scalar and vector fields, implemented using just-in-time compilation
jitfields.readthedocs.io
Resources
Readme
MIT license
Activity
Stars
2 stars
Watchers
2 watching
Forks
2 forks
Report repository
Releases
Packages
Used by
Contributors
Languages