[FIX] fix progressive affine when an initial guess is provided

Author: balbasty

nitorch/tools/registration/objects.py

@@ -2043,27 +2043,24 @@ def __init__(self, model: "AffineModel"):
         optim = property(lambda self: self.model.optim)
         basis = property(lambda self: self.model.basis_name)
 
-        def add_(self, value, **kwargs):
+        def _switch(self, value):
             if self.optim not in (None, self.basis):
                 value = LogAffine._switch_basis(value, self.optim, self.basis, self.model.dat.dim)
-            self.dat.add_(value, **kwargs)
+            return value
+
+        def add_(self, value, **kwargs):
+            self.dat.add_(self._switch(value), **kwargs)
             return self
 
         def add(self, value, **kwargs):
-            if self.optim not in (None, self.basis):
-                value = LogAffine._switch_basis(value, self.optim, self.basis, self.model.dat.dim)
-            return self.dat.add(value, **kwargs)
+            return self.dat.add(self._switch(value), **kwargs)
 
         def mul_(self, value, **kwargs):
-            if self.optim not in (None, self.basis):
-                value = LogAffine._switch_basis(value, self.optim, self.basis, self.model.dat.dim)
-            self.dat.mul_(value, **kwargs)
+            self.dat.mul_(self._switch(value), **kwargs)
             return self
 
         def mul(self, value, **kwargs):
-            if self.optim not in (None, self.basis):
-                value = LogAffine._switch_basis(value, self.optim, self.basis, self.model.dat.dim)
-            return self.dat.mul(value, **kwargs)
+            return self.dat.mul(self._switch(value), **kwargs)
 
         def __add__(self, other):
             return self.add(other)
@@ -2247,6 +2244,8 @@ def __repr__(self):
         else:
             s += ['<uninitialized>']
         s += [f'basis={self._basis}']
+        if self.optim is not None:
+            s += [f'optim={self.optim}']
         s += [f'factor={self.factor}']
         s += [f'position={self.position}']
         if self.penalty is not None:
@@ -2260,27 +2259,10 @@ class Affine2dModel(AffineModel):
 
     class Parameters(AffineModel.Parameters):
 
-        def add_(self, value, **kwargs):
-            if self.optim not in (None, self.basis):
-                value = LogAffine2d._switch_basis(value, self.optim, self.basis)
-            self.dat.add_(value, **kwargs)
-            return self
-
-        def add(self, value, **kwargs):
-            if self.optim not in (None, self.basis):
-                value = LogAffine._switch_basis(value, self.optim, self.basis, self.model.dat.dim)
-            return self.dat.add(value, **kwargs)
-
-        def mul_(self, value, **kwargs):
+        def _switch(self, value):
             if self.optim not in (None, self.basis):
                 value = LogAffine2d._switch_basis(value, self.optim, self.basis)
-            self.dat.mul_(value, **kwargs)
-            return self
-
-        def mul(self, value, **kwargs):
-            if self.optim not in (None, self.basis):
-                value = LogAffine._switch_basis(value, self.optim, self.basis, self.model.dat.dim)
-            return self.dat.mul(value, **kwargs)
+            return value
 
     def __init__(self, basis, plane, ref_affine=None, factor=1, penalty=None,
                  dat=None, position='symmetric'):
@@ -2307,7 +2289,13 @@ def plane_name(self):
         return self._plane
 
     def set_dat(self, dat=None, dim=None, **backend):
-        self.dat = LogAffine2d(dat, basis=self._basis, dim=dim,
-                               rotation=self._rotation,
-                               plane=self._plane, **backend)
+        self.dat = LogAffine2d(
+            dat,
+            basis=self._basis,
+            optim=self._optim,
+            dim=dim,
+            rotation=self._rotation,
+            plane=self._plane,
+            **backend
+        )
         return self

nitorch/tools/registration/pairwise_run.py

@@ -106,8 +106,7 @@ def run(
     # --- progressive affine initialization ---
     if affine and not affine.frozen and progressive:
         affine, figure = run_progressive_init(
-            runner, losses, affine, affine_optim, line_size=line_size,
-            **plotopt)
+            runner, losses, affine, affine_optim, **plotopt)
         plotopt["figure"] = figure
 
     # --- full affine ---
@@ -185,16 +184,16 @@ def run_progressive_init(
         return affine, None
 
     if verbose:
-        print('-' * line_size)
+        print('=' * line_size)
         print('   PROGRESSIVE INITIALIZATION')
-        print('-' * line_size)
+        print('=' * line_size)
 
     affine.optim = names.pop(0)
     while names:
         if verbose:
             line_pad = line_size - len(affine.optim) - 5
-            print(f'--- {affine.optim} ', end='')
-            print('-' * max(0, line_pad))
+            print(f'*** {affine.optim} ', end='')
+            print('*' * max(0, line_pad))
         affine_optim.reset_state()
         torch.cuda.empty_cache()
 
@@ -244,11 +243,21 @@ def run_pyramid(
     optim,
     verbose=True,
     framerate=1,
-    figure=None
+    figure=None,
 ):
     """Run sequential pyramid registration"""
     line_size = 89 if nonlin else 74
 
+    if verbose:
+        print('=' * line_size)
+        if affine and not nonlin:
+            print('   AFFINE')
+        elif nonlin and not affine:
+            print('   NONLINEAR')
+        else:
+            print('   AFFINE & NONLINEAR')
+        print('=' * line_size)
+
     nonlin_optim = None
     if nonlin and not getattr(nonlin, 'frozen', False):
         if affine and not getattr(affine, 'frozen', False):