flows

This module contains finite-difference regularisers and associated utilities for tensors that contain dense displacement fields.

Displacements are expected to be expressed in voxels, and relate to the tensor's lattice. However, the voxel size of the lattice (and therefore, the size of the displacement) can be provided.

Each function implement the following set of penalties, which can be combined:

absolute: the absolute energy is the sum of squared values, i.e., \(\mathcal{L} = \frac{1}{2} \int \lVert \boldsymbol{f}\left(\mathbf{x}\right) \rVert_2^2 \mathrm{d}\mathbf{x}\).
membrane: the membrane energy is the sum of squared first order derivatives, i.e., \(\mathcal{L} = \frac{1}{2} \int \lVert \left(\mathbf{D}\boldsymbol{f}\right)\left(\mathbf{x}\right) \rVert_F^2 \mathrm{d}\mathbf{x}\), where \(\mathbf{D}\boldsymbol{f}\) is the Jacobian matrix of the displacement \(\boldsymbol{f}\).
bending: the bending energy is the sum of squared second order derivatives, i.e., \(\mathcal{L} = \frac{1}{2} \int \sum_d \lVert \left(\mathbf{H}f_d\right)\left(\mathbf{x}\right) \rVert_F^2 \mathrm{d}\mathbf{x}\), where \(\mathbf{H}f_d\) is the Hessian matrix of the d-th component of the displacement \(\boldsymbol{f}\).
shears: the second Lame component of the linear-elastic energy penalizes the norm of the symmetric part of the Jacobian, i.e., \(\mathcal{L} = \frac{1}{2} \int \lVert \frac{1}{2} \left(\mathbf{D}\boldsymbol{f} + \mathbf{D}\boldsymbol{f}^{T}\right) \left(\mathbf{x}\right) \rVert_F^2 \mathrm{d}\mathbf{x}\).
div: the first Lame component of the linear-elastic energy penalizes divergence, which is the square of the trace of the Jacobian, i.e. \(\mathcal{L} = \frac{1}{2} \int \operatorname{tr}\left(\mathbf{D}\boldsymbol{f}\right)^2 \left(\mathbf{x}\right) \mathrm{d}\mathbf{x}\).

We further allow a local weighting of the penalty, which in the case of the membrane energy yields \(\mathcal{L} = \frac{1}{2} \int w\left(\mathbf{x}\right) \lVert \left(\mathbf{D}\boldsymbol{f}\right)\left(\mathbf{x}\right) \rVert_F^2 \mathrm{d}\mathbf{x}\).

In practice, \(\boldsymbol{f}\) is a dense discrete field, i.e., \(\mathbf{f} \in \mathbb{R}^{ND}\) (where \(N\) is the number of voxels). All these penalties are then quadratic forms in \(\mathbf{f}\): \(\mathcal{L} = \frac{1}{2}\mathbf{f}^{T}\mathbf{Lf}\). The penalties are implemented using finite difference and, in the absence of local weighting, \(\mathbf{L}\) takes the form of a convolution with a small kernel.

flow_matvec

flow_matvec(flow, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None)

Apply a spatial regularization matrix.

Note

This function computes the matrix-vector product \(\mathbf{L} \times \mathbf{f}\), where \(\mathbf{f}\) is a displacement field (in voxels) and \(\mathbf{L}\) encodes a finite-difference penalty.

Parameters:

Name	Type	Description	Default
`flow`	`(batch, spatial, ndim) tensor`	Input displacement field, in voxels. With shape `(batch, spatial, ndim)`.	required
`weight`	`(batch, spatial) tensor`	Weight map, to spatially modulate the regularization. With shape `(batch, spatial)`.	`None`
`absolute`	`float`	Penalty on absolute values.	`0`
`membrane`	`float`	Penalty on first derivatives.	`0`
`bending`	`float`	Penalty on second derivatives.	`0`
`shears`	`float`	Penalty on local shears.	`0`
`div`	`float`	Penalty on local volume changes.	`0`
`bound`	`[sequence of] BoundLike`	Boundary conditions.	`'dft'`
`voxel_size`	`[sequence of] float`	Voxel size.	`1`
`out`	`(batch, spatial, ndim) tensor`	Output placeholder. With shape `(batch, spatial, ndim)`.	`None`

Returns:

Name	Type	Description
`out`	`(batch, spatial, ndim) tensor`	Matrix vector product, with shape `(batch, spatial, ndim)`.

flow_matvec_add

flow_matvec_add(inp, flow, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None, _sub=False)

Add the output of flow_matvec to inp.

flow_matvec_add_

flow_matvec_add_(inp, flow, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, _sub=False)

Add the output of flow_matvec to inp (inplace).

flow_matvec_sub

flow_matvec_sub(inp, flow, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None)

Subtract the output of flow_matvec from inp.

flow_matvec_sub_

flow_matvec_sub_(inp, flow, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1)

Subtract the output of flow_matvec from inp (inplace).

flow_kernel

flow_kernel(shape, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None, **backend)

Return the kernel of a Toeplitz regularization matrix.

Parameters:

Name	Type	Description	Default
`shape`	`int or list[int]`	Number of spatial dimensions or shape of the tensor.	required
`absolute`	`float`	Penalty on absolute values.	`0`
`membrane`	`float`	Penalty on first derivatives.	`0`
`bending`	`float`	Penalty on second derivatives.	`0`
`shears`	`float`	Penalty on local shears. Linear elastic energy's `mu`.	`0`
`div`	`float`	Penalty on local volume changes. Linear elastic energy's `lambda`.	`0`
`bound`	`[sequence of] BoundLike`	Boundary conditions.	`'dft'`
`voxel_size`	`[sequence of] float`	Voxel size.	`1`
`out`	`(*shape, ndim, [ndim]) tensor`	Output placeholder, with shape `(*shape, ndim, [ndim])`.	`None`

Returns:

Name	Type	Description
`out`	`(*shape, ndim, [ndim]) tensor`	Convolution kernel. A matrix or kernels with shape `(shape, ndim, ndim)` if `shears` or `div`, else a vector of kernels with shape `(shape, ndim)`.

flow_kernel_add

flow_kernel_add(inp, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None, _sub=False)

Add the output of flow_kernel to inp.

flow_kernel_add_

flow_kernel_add_(inp, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, _sub=False)

Add the output of flow_kernel to inp (inplace).

flow_kernel_sub

flow_kernel_sub(inp, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None)

Subtract the output of flow_kernel from inp.

flow_kernel_sub_

flow_kernel_sub_(inp, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1)

Subtract the output of flow_kernel from inp (inplace).

flow_diag

flow_diag(shape, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None, **backend)

Return the diagonal of a regularization matrix.

Parameters:

Name	Type	Description	Default
`shape`	`list[int]`	Shape of the tensor.	required
`weight`	`(batch, spatial) tensor`	Weight map, to spatially modulate the regularization. With shape `(batch, spatial)`.	`None`
`absolute`	`float`	Penalty on absolute values.	`0`
`membrane`	`float`	Penalty on first derivatives.	`0`
`bending`	`float`	Penalty on second derivatives.	`0`
`shears`	`float`	Penalty on local shears.	`0`
`div`	`float`	Penalty on local volume changes.	`0`
`bound`	`[sequence of] BoundLike`	Boundary conditions.	`'dft'`
`voxel_size`	`[sequence of] float`	Voxel size.	`1`
`out`	`(batch, spatial, ndim) tensor`	Output placeholder	`None`

Returns:

Name	Type	Description
`out`	`(batch, spatial, ndim) tensor`	Diagonal of the regularisation matrix, with shape `(batch, spatial, ndim)`.

flow_diag_add

flow_diag_add(inp, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None, _sub=False)

Add the output of flow_diag to inp.

flow_diag_add_

flow_diag_add_(inp, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, _sub=False)

Add the output of flow_diag to inp (inplace).

flow_diag_sub

flow_diag_sub(inp, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None)

Subtract the output of flow_diag from inp.

flow_diag_sub_

flow_diag_sub_(inp, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1)

Subtract the output of flow_diag from inp (inplace).

flow_precond

flow_precond(mat, vec, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None)

Apply the preconditioning (M + diag(R)) \ v

Parameters:

Name	Type	Description	Default
`mat`	`(batch, spatial, DD) tensor`	Preconditioning matrix `M` with shape `(batch, spatial, DD)`, where `DD` is one of `{1, D, D(D+1)//2, DD}`.	required
`vec`	`(batch, spatial, D) tensor`	Point `v` at which to solve the system, with shape `(batch, spatial, D)`.	required
`weight`	`(`batch, spatial) tensor`	Regularization weight map, with shape `(batch, spatial)`.	`None`
`absolute`	`float`	Penalty on absolute values.	`0`
`membrane`	`float`	Penalty on first derivatives.	`0`
`bending`	`float`	Penalty on second derivatives.	`0`
`shears`	`float`	Penalty on local shears.	`0`
`div`	`float`	Penalty on local volume changes.	`0`
`bound`	`[sequence of] BoundLike`	Boundary conditions.	`'dft'`
`voxel_size`	`[sequence of] float`	Voxel size.	`1`
`out`	`(batch, spatial, D) tensor`	Output placeholder, with shape	`None`

Returns:

Name	Type	Description
`out`	`(batch, spatial, D) tensor`	Preconditioned vector.

flow_precond_

flow_precond_(mat, vec, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1)

Apply the preconditioning (M + diag(R)) \ v inplace. See flow_precond.

flow_forward

flow_forward(mat, vec, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, out=None)

Apply the forward matrix-vector product (M + R) @ v

Parameters:

Name	Type	Description	Default
`mat`	`(batch, spatial, DD) tensor`	Block-diagonal matrix with shape `(batch, spatial, DD)`, where `DD` is one of `{1, D, D(D+1)//2, DD}`.	required
`vec`	`(batch, spatial, D) tensor`	Point `v` at which to solve the system, with shape `(batch, spatial, D)`	required
`weight`	`(batch, spatial) tensor`	Regularization weight map, with shape `(batch, spatial)`.	`None`
`absolute`	`float`	Penalty on absolute values.	`0`
`membrane`	`float`	Penalty on first derivatives.	`0`
`bending`	`float`	Penalty on second derivatives.	`0`
`shears`	`float`	Penalty on local shears.	`0`
`div`	`float`	Penalty on local volume changes.	`0`
`bound`	`[sequence of] BoundLike`	Boundary conditions.	`'dft'`
`voxel_size`	`[sequence of] float`	Voxel size.	`1`
`out`	`(batch, spatial, D) tensor`	Output placeholder, with shape `(batch, spatial, D)`.	`None`

Returns:

Name	Type	Description
`out`	`(batch, spatial, D) tensor`	Preconditioned vector, with shape `(batch, spatial, D)`.

flow_relax_

flow_relax_(flow, hes, grd, weight=None, absolute=0, membrane=0, bending=0, shears=0, div=0, bound='dft', voxel_size=1, nb_iter=1)

Perform relaxation iterations (inplace).

Parameters:

Name	Type	Description	Default
`flow`	`(batch, spatial, ndim) tensor`	Warm start, with shape `(batch, spatial, ndim)`.	required
`hes`	`(batch, spatial, ndim*(ndim+1)//2) tensor`	Input symmetric Hessian, in voxels. With shape `(batch, spatial, ndim*(ndim+1)//2)`.	required
`grd`	`(batch, spatial, ndim) tensor`	Input gradient, in voxels. With shape `(batch, spatial, ndim)`.	required
`weight`	`(batch, spatial) tensor`	Weight map, to spatially modulate the regularization. With shape `(batch, spatial)`.	`None`
`absolute`	`float`	Penalty on absolute values.	`0`
`membrane`	`float`	Penalty on first derivatives.	`0`
`bending`	`float`	Penalty on second derivatives.	`0`
`shears`	`float`	Penalty on local shears.	`0`
`div`	`float`	Penalty on local volume changes.	`0`
`bound`	`[sequence of] BoundLike`	Boundary conditions.	`'dft'`
`voxel_size`	`[sequence of] float`	Voxel size.	`1`
`nb_iter`	`int`	Number of iterations	`1`

Returns:

Name	Type	Description
`flow`	`(batch, spatial, ndim) tensor`	Refined solution with shape `(batch, spatial, ndim)` .