ENH(register): better partial affine optimization

Author: balbasty

nitorch/cli/registration/register/cli.py

@@ -278,7 +278,8 @@ def build_losses(options, pyramids, device):
         image_dict[loss.mov.name or loss.mov.files[0]] = mov
 
         # Forward loss
-        factor = loss.factor / (2 if loss.symmetric else 1)
+        loss.factor = py.ensure_list(loss.factor, len(pyramid))[::-1]
+        factor = loss.factor[-1] / (2 if loss.symmetric else 1)
         sumloss = sumloss + factor * objects.Similarity(lossobj, mov, fix)
 
         # Backward loss
@@ -296,6 +297,40 @@ def build_losses(options, pyramids, device):
     pyramid_fn = (concurrent_pyramid if options.pyramid.concurrent else
                   sequential_pyramid)
     loss_list = pyramid_fn(sumloss)
+
+    if options.pyramid.concurrent:
+        if not isinstance(loss_list, objects.SumSimilarity):
+            loss_list = objects.SumSimilarity([loss_list])
+        loss_list = list(loss_list)
+        new_list = []
+        nb_levels = len(loss_list) // len(options.loss)
+        nb_losses = len(options.loss)
+        for i in range(nb_levels):
+            for j in range(nb_losses):
+                rel_factor = (
+                    options.loss[j].factor[i] /
+                    options.loss[j].factor[-1]
+                )
+                new_list.append(loss_list[i*nb_losses+j] * rel_factor)
+        loss_list = objects.SumSimilarity.sum(new_list)
+    else:
+        new_list = []
+        for i, level in enumerate(loss_list):
+            if not isinstance(level, objects.SumSimilarity):
+                level = objects.SumSimilarity([level])
+            level = list(level)
+            new_level = []
+            for j, loss_elem in enumerate(level):
+                i = min(i, len(options.loss[j].factor) - 1)
+                rel_factor = (
+                    options.loss[j].factor[i] /
+                    options.loss[j].factor[-1]
+                )
+                new_level.append(loss_elem * rel_factor)
+            new_level = objects.SumSimilarity.sum(new_level)
+            new_list.append(new_level)
+        loss_list = new_list
+
     return loss_list, image_dict
 
 

nitorch/cli/registration/register/parser.py

@@ -95,7 +95,6 @@
                     @pyramid [NAME] ...
 
 @loss options:
-    FACTOR must be a scalar value [1]
     NAME can take values (with specific sub-options):
        [mi, nmi]                        Mutual information (can be normalized)
             -m, --norm NAME                 Normalization: [studholme], arithmetic, geometric, no
@@ -126,6 +125,7 @@
         cat, cce                        Categorical cross-entropy
         dice, f1                        Dice coefficient
             -w, --weight *VAL               Weight per class [1]
+    FACTOR must be a scalar value [1]
     Common options:
         -s, --symmetric                 Make loss symmetric [False]
         -z, --slicewise [AXIS=-1]       Make loss slice-wise [False]
@@ -153,13 +153,13 @@
     -c, --channels *C               Channels to load. Can be a range start:stop:step [:]
 
 @affine options:
-    FACTOR must be a scalar value [1] and is a global penalty factor
     NAME can take values:
         t, translation                  Translations only
         o, rotation                     Rotations only
        [r, rigid]                       Translations + Rotations
         s, similitude                   Translations + Rotations + Iso zoom
         a, affine                       Full affine
+    FACTOR must be a scalar value [1] and is a global penalty factor
     Common options:
         -p, --position                  Position of the affine: [sym], mov, fix
         -g, --progressive               Progressive optimization (t -> r -> s -> a) [false]
@@ -168,11 +168,11 @@
         -2d [AXIS=2]                    Force transform to be 2d about AXIS
 
 @nonlin options:
-    FACTOR must be a scalar value [1] and is a global penalty factor
     NAME can take values:
        [v, svf]                         Stationary velocity field
         g, shoot                        Geodesic shooting
         d, smalldef                     Dense deformation field
+    FACTOR must be a scalar value [1] and is a global penalty factor
     Common options:
         -i, --input                     Path to initial transform
         -o, --output                    Path to the output transform: [{dir}/{name}.nii.gz]
@@ -336,7 +336,7 @@ def _convert(x):
                               help='Name of the image loss')
 loss = cli.NamedGroup('loss', loss_choices, '@loss', n='+',
                       help='A loss between two images')
-loss.add_positional('factor', nargs='?', default=1., convert=float,
+loss.add_positional('factor', nargs='*', default=[1.], convert=float,
                     help='Weight it this component in the global loss')
 loss.add_option('symmetric', ('-s', '--symmetric'), nargs=0, default=False,
                 help='Make the loss symmetric')
@@ -529,7 +529,7 @@ def _convert(x):
                                 convert=lambda x: affine_aliases.get(x, x),
                                 help='Name of the affine transform')
 affine = cli.NamedGroup('affine', affine_choices, '@affine', n='?', make_default=False)
-affine.add_positional('factor', nargs='?', default=1., convert=float,
+affine.add_positional('factor', nargs='*', default=[1.], convert=float,
                       help='Penalty factor')
 affine.add_option('position', ('-p', '--position'), default='sym', nargs=1,
                   validation=cli.Validations.choice(['sym', 'mov', 'fix']),
@@ -556,7 +556,7 @@ def _convert(x):
                                 convert=lambda x: nonlin_aliases.get(x, x),
                                 help='Name of the nonlinear transform')
 nonlin = cli.NamedGroup('nonlin', nonlin_choices, '@nonlin', n='?', make_default=False)
-nonlin.add_positional('factor', nargs='?', default=1., convert=float,
+nonlin.add_positional('factor', nargs='*', default=[1.], convert=float,
                       help='Penalty factor')
 nonlin.add_option('steps', ('-s', '--steps'), nargs=1, default=8, convert=int,
                   help='Number of integration steps')