fiery.diffeo.metrics

Riemannian metrics on velocity fields.

A metric is a positive semi-definite linear operator L that maps a velocity field to its momentum field. It plays two roles in this package: it regularizes velocity fields (the penalty is the inner product (v, Lv)), and it defines the geodesics that fiery.diffeo.layers.Shoot integrates.

Three families are available:

Mixture builds L from a weighted sum of finite-difference energies (absolute, membrane, bending and linear-elastic);
Laplace and Helmoltz use the analytical Green's function of the corresponding differential operator;
Gaussian uses a Gaussian filter as its Green's function.

All of them derive from the Metric base class and therefore share the same forward / inverse / whiten / color / logdet interface.

Classes

Laplace

Laplace(factor=1, voxel_size=1, bound='circulant', learnable=False, cache=False)

Bases: Metric

Positive semi-definite metric based on the Laplace operator.

This is relatively similar to SPM's "membrane" energy, but relies on the (ill-posed) analytical form of the Green's function.

See:

Parameters:

Name	Type	Description	Default
`factor`	`float`	Regularization factor (optionally: learnable).	`1`
`voxel_size`	`[list of] float`	Voxel size.	`1`
`bound`	`[list of] {'circulant', 'neumann', 'dirichlet', 'sliding'}`	Boundary conditions.	`'circulant'`
`learnable`	`bool`	Make `factor` a learnable parameter.	`False`
`cache`	`bool or int`	Cache up to `n` kernels.	`False`

Methods:

forward

forward(x, factor=True)

Apply the forward linear operator: v -> Lv.

Converts a velocity field into a momentum field.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`m`	`(..., *spatial, D) tensor`	Output momentum field.

inverse

inverse(x, factor=True)

Apply the inverse (Green's) linear operator: m -> Km.

Converts a momentum field back into a velocity field, where K = inv(L).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input momentum field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

whiten

whiten(x, factor=True)

Apply the square root of the forward operator: v -> sqrt(L) v.

Whitens a velocity field: under the Gaussian prior whose precision matrix is L, the output has an identity covariance.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`w`	`(..., *spatial, D) tensor`	Output white field.

color

color(x, factor=True)

Apply the square root of the inverse operator: w -> sqrt(K) w.

Colors a white field: the output is a velocity field distributed according to the Gaussian prior whose precision matrix is L. This is the inverse of whiten.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input white field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

logdet

logdet(x, factor=True)

Return the log-determinant of L (times the batch size).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	A velocity field. Only its shape, dtype and device are used; its values are not.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`ld`	`scalar tensor`	Log-determinant, scaled by the batch size.

Helmoltz

Helmoltz(factor=1, alpha=0.001, voxel_size=1, bound='circulant', learnable=False, cache=False)

Bases: Metric

Positive semi-definite metric based on the Helmholtz operator.

This is relatively similar to SPM's mixture of "absolute" and "membrane" energies, but relies on the (ill-posed) analytical form of the Green's function.

See:

Parameters:

Name	Type	Description	Default
`factor`	`float`	Regularization factor (optionally: learnable).	`1`
`alpha`	`float`	Diagonal regularizer (cannot be learned). It is the square of the eigenvalue in the Helmholtz equation.	`0.001`
`voxel_size`	`[list of] float`	Voxel size.	`1`
`bound`	`[list of] {'circulant', 'neumann', 'dirichlet', 'sliding'}`	Boundary conditions.	`'circulant'`
`learnable`	`bool`	Make `factor` a learnable parameter.	`False`
`cache`	`bool or int`	Cache up to `n` kernels.	`False`

Methods:

forward

forward(x, factor=True)

Apply the forward linear operator: v -> Lv.

Converts a velocity field into a momentum field.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`m`	`(..., *spatial, D) tensor`	Output momentum field.

inverse

inverse(x, factor=True)

Apply the inverse (Green's) linear operator: m -> Km.

Converts a momentum field back into a velocity field, where K = inv(L).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input momentum field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

whiten

whiten(x, factor=True)

Apply the square root of the forward operator: v -> sqrt(L) v.

Whitens a velocity field: under the Gaussian prior whose precision matrix is L, the output has an identity covariance.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`w`	`(..., *spatial, D) tensor`	Output white field.

color

color(x, factor=True)

Apply the square root of the inverse operator: w -> sqrt(K) w.

Colors a white field: the output is a velocity field distributed according to the Gaussian prior whose precision matrix is L. This is the inverse of whiten.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input white field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

logdet

logdet(x, factor=True)

Return the log-determinant of L (times the batch size).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	A velocity field. Only its shape, dtype and device are used; its values are not.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`ld`	`scalar tensor`	Log-determinant, scaled by the batch size.

Mixture

Mixture(absolute=0, membrane=0, bending=0, lame_shears=0, lame_div=0, factor=1, voxel_size=1, bound='circulant', use_diff=True, learnable=False, cache=False)

Bases: Metric

Positive semi-definite metric based on finite-difference regularisers.

Mixture of "absolute", "membrane", "bending" and "linear-elastic" energies. Note that these names refer to what is penalised when computing the inner product (v, Lv). The "membrane" energy is therefore closely related to the "Laplacian" metric.

Parameters:

Name	Type	Description	Default
`absolute`	`float`	Penalty on (squared) absolute values.	`0`
`membrane`	`float`	Penalty on (squared) first derivatives.	`0`
`bending`	`float`	Penalty on (squared) second derivatives.	`0`
`lame_shears`	`float`	Penalty on the (squared) symmetric component of the Jacobian.	`0`
`lame_div`	`float`	Penalty on the trace of the Jacobian.	`0`
`factor`	`float`	Global regularization factor (optionally: learnable).	`1`
`voxel_size`	`[list of] float`	Voxel size.	`1`
`bound`	`[list of] {'circulant', 'neumann', 'dirichlet', 'sliding'}`	Boundary conditions.	`'circulant'`
`use_diff`	`bool`	Use finite differences to perform the forward pass. Otherwise, perform the convolution in Fourier space.	`True`
`learnable`	`bool or {factor, components, factor + components}`	Make `factor` a learnable parameter. If `'components'`, the individual factors (absolute, membrane, etc.) are learned instead of the global factor. Individual components can also be named one by one, e.g. `'membrane+bending'`. `True` is equivalent to `'factor'`.	`False`
`cache`	`bool or int`	Cache up to `n` kernels. This cannot be used together with learnable components.	`False`

Methods:

inverse

inverse(x, factor=True)

Apply the inverse (Green's) linear operator: m -> Km.

Converts a momentum field back into a velocity field, where K = inv(L).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input momentum field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

whiten

whiten(x, factor=True)

Apply the square root of the forward operator: v -> sqrt(L) v.

Whitens a velocity field: under the Gaussian prior whose precision matrix is L, the output has an identity covariance.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`w`	`(..., *spatial, D) tensor`	Output white field.

color

color(x, factor=True)

Apply the square root of the inverse operator: w -> sqrt(K) w.

Colors a white field: the output is a velocity field distributed according to the Gaussian prior whose precision matrix is L. This is the inverse of whiten.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input white field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

logdet

logdet(x, factor=True)

Return the log-determinant of L (times the batch size).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	A velocity field. Only its shape, dtype and device are used; its values are not.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`ld`	`scalar tensor`	Log-determinant, scaled by the batch size.

forward

forward(x, factor=True)

Apply the forward linear operator: v -> Lv.

Converts a velocity field into a momentum field. When use_diff is True this is computed directly with finite differences; otherwise it is computed as a convolution in Fourier space.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`m`	`(..., *spatial, D) tensor`	Output momentum field.

Gaussian

Gaussian(fwhm=16, factor=1, voxel_size=1, bound='circulant', learnable=False, cache=False)

Bases: Metric

Positive semi-definite metric whose Green's function is a Gaussian filter.

Parameters:

Name	Type	Description	Default
`fwhm`	`float`	Full-width at half-maximum of the Gaussian filter, in mm (optionally: learnable).	`16`
`factor`	`float`	Global regularization factor (optionally: learnable).	`1`
`voxel_size`	`[list of] float`	Voxel size.	`1`
`bound`	`[list of] {'circulant', 'neumann', 'dirichlet', 'sliding'}`	Boundary conditions.	`'circulant'`
`learnable`	`bool or {factor, fwhm, fwhm + factor}`	Make `factor` and/or `fwhm` a learnable parameter. `True` is equivalent to `'factor'`.	`False`
`cache`	`bool or int`	Cache up to `n` kernels. This cannot be used together with a learnable `fwhm`.	`False`

Raises:

Type	Description
`ValueError`	If `cache` is used together with a learnable `fwhm`.

Methods:

forward

forward(x, factor=True)

Apply the forward linear operator: v -> Lv.

Converts a velocity field into a momentum field.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`m`	`(..., *spatial, D) tensor`	Output momentum field.

inverse

inverse(x, factor=True)

Apply the inverse (Green's) linear operator: m -> Km.

Converts a momentum field back into a velocity field, where K = inv(L).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input momentum field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

whiten

whiten(x, factor=True)

Apply the square root of the forward operator: v -> sqrt(L) v.

Whitens a velocity field: under the Gaussian prior whose precision matrix is L, the output has an identity covariance.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input velocity field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`w`	`(..., *spatial, D) tensor`	Output white field.

color

color(x, factor=True)

Apply the square root of the inverse operator: w -> sqrt(K) w.

Colors a white field: the output is a velocity field distributed according to the Gaussian prior whose precision matrix is L. This is the inverse of whiten.

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	Input white field.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`v`	`(..., *spatial, D) tensor`	Output velocity field.

logdet

logdet(x, factor=True)

Return the log-determinant of L (times the batch size).

Parameters:

Name	Type	Description	Default
`x`	`(..., *spatial, D) tensor`	A velocity field. Only its shape, dtype and device are used; its values are not.	required
`factor`	`bool`	Whether to incorporate the global regularization factor.	`True`

Returns:

Name	Type	Description
`ld`	`scalar tensor`	Log-determinant, scaled by the batch size.