def __init__(self, dim, basis='CSO', fwd=True, inv=False):
"""
Parameters
----------
dim : {1, 2, 3}
Spatial dimension
basis : basis_like or list[basis_like], default='CSO'
The simplest way to define an affine basis is to choose from
a list of Lie groups:
* 'T' : Translations
* 'SO' : Special Orthogonal (rotations)
* 'SE' : Special Euclidean (translations + rotations)
* 'D' : Dilations (translations + isotropic scalings)
* 'CSO' : Conformal Special Orthogonal
(translations + rotations + isotropic scalings)
* 'SL' : Special Linear (rotations + isovolumic zooms + shears)
* 'GL+' : General Linear [det>0] (rotations + zooms + shears)
* 'Aff+': Affine [det>0] (translations + rotations + zooms + shears)
More complex (hierarchical) encodings can be achieved as well.
See `affine_matrix`.
fwd : bool, default=True
Return the forward transformation.
inv : bool, default=False
Return the inverse transformation.
"""
super().__init__()
self.dim = dim
self.basis = spatial.build_affine_basis(basis, dim)
self.fwd = fwd
self.inv = inv
def forward(self, prm, **overload):
"""
Parameters
----------
prm : (batch, nb_prm) tensor or list[tensor]
Affine parameters on the Lie algebra.
overload : dict
All parameters of the module can be overridden at call time.
Returns
-------
forward : (batch, dim+1, dim+1) tensor, optional
Forward matrix
inverse : (batch, dim+1, dim+1) tensor, optional
Inverse matrix
"""
fwd = overload.get('fwd', self.forward)
inv = overload.get('inverse', self.inv)
basis = overload.get('basis', self.basis)
basis = spatial.build_affine_basis(basis, self.dim)
output = []
if fwd:
aff = spatial.affine_matrix(prm, basis)
output.append(aff)
if inv:
if isinstance(prm, (list, tuple)):
prm = [-p for p in prm]
else:
prm = -prm
iaff = spatial.affine_matrix(prm, basis)
output.append(iaff)
return output if len(output) > 1 else \
output[0] if len(output) == 1 else \
None
def forward(self, affine, **overload):
"""
Parameters
----------
prm : (batch, nb_prm) tensor or list[tensor]
Affine parameters on the Lie algebra.
overload : dict
All parameters of the module can be overridden at call time.
Returns
-------
forward : (batch, dim+1, dim+1) tensor, optional
Forward matrix
inverse : (batch, dim+1, dim+1) tensor, optional
Inverse matrix
"""
# I build the matrix at each call, which is not great.
# Hard to be efficient *and* generic...
dim = affine.shape[-1] - 1
backend = dict(dtype=affine.dtype, device=affine.device)
basis = overload.get('basis', self.basis)
basis = spatial.build_affine_basis(basis, dim, **backend)
# When the affine is well conditioned, its log should be real.
# Here, I take the real part just in case.
# Another solution could be to regularise the affine (by loading
# slightly the diagonal) until it is well conditioned -- but
# how would that work with autograd?
affine = core.linalg.logm(affine.double())
if affine.is_complex():
affine = affine.real
affine = affine.to(**backend)
affine = core.linalg.mdot(affine[:, None, ...], basis[None, ...])
return affine
def __init__(self, dim, basis='CSO', **backend):
"""
Parameters
----------
dim : int
Number of spatial dimensions
basis : basis_like or list[basis_like], default='CSO'
The simplest way to define an affine basis is to choose from
a list of Lie groups:
* 'T' : Translations
* 'SO' : Special Orthogonal (rotations)
* 'SE' : Special Euclidean (translations + rotations)
* 'D' : Dilations (translations + isotropic scalings)
* 'CSO' : Conformal Special Orthogonal
(translations + rotations + isotropic scalings)
* 'SL' : Special Linear (rotations + isovolumic zooms + shears)
* 'GL+' : General Linear [det>0] (rotations + zooms + shears)
* 'Aff+': Affine [det>0] (translations + rotations + zooms + shears)
More complex (hierarchical) encodings can be achieved as well.
See `affine_matrix`.
"""
super().__init__()
self.basis = spatial.build_affine_basis(basis, dim, **backend)