[ENH/FIX] Changes for reward bundles and VOF

.codespellrc

@@ -1,3 +1,3 @@
 [codespell]
 skip = [setup.cfg] 
-ignore-words-list = Reson, DNE, ACI, FPT, sagital, saggital, abd, Joo, Mapp, Commun, vor, Claus
+ignore-words-list = Reson, DNE, ACI, FPT, sagital, saggital, abd, Joo, Mapp, Commun, vor, Claus, coo

AFQ/_fixes.py

@@ -4,6 +4,8 @@
 from math import radians
 
 import numpy as np
+from dipy.align import vector_fields as vfu
+from dipy.align.imwarp import DiffeomorphicMap, mult_aff
 from dipy.data import default_sphere
 from dipy.reconst.gqi import squared_radial_component
 from dipy.tracking.streamline import set_number_of_points
@@ -15,6 +17,75 @@
 logger = logging.getLogger("AFQ")
 
 
+def get_simplified_transform(self):
+    """Constructs a simplified version of this Diffeomorhic Map
+
+    The simplified version incorporates the pre-align transform, as well as
+    the domain and codomain affine transforms into the displacement field.
+    The resulting transformation may be regarded as operating on the
+    image spaces given by the domain and codomain discretization. As a
+    result, self.prealign, self.disp_grid2world, self.domain_grid2world and
+    self.codomain affine will be None (denoting Identity) in the resulting
+    diffeomorphic map.
+    """
+    if self.dim == 2:
+        simplify_f = vfu.simplify_warp_function_2d
+    else:
+        simplify_f = vfu.simplify_warp_function_3d
+    # Simplify the forward transform
+    D = self.domain_grid2world
+    P = self.prealign
+    Rinv = self.disp_world2grid
+    Cinv = self.codomain_world2grid
+
+    # this is the matrix which we need to multiply the voxel coordinates
+    # to interpolate on the forward displacement field ("in"side the
+    # 'forward' brackets in the expression above)
+    affine_idx_in = mult_aff(Rinv, mult_aff(P, D))
+
+    # this is the matrix which we need to multiply the voxel coordinates
+    # to add to the displacement ("out"side the 'forward' brackets in the
+    # expression above)
+    affine_idx_out = mult_aff(Cinv, mult_aff(P, D))
+
+    # this is the matrix which we need to multiply the displacement vector
+    # prior to adding to the transformed input point
+    affine_disp = Cinv
+
+    new_forward = simplify_f(
+        self.forward, affine_idx_in, affine_idx_out, affine_disp, self.domain_shape
+    )
+
+    # Simplify the backward transform
+    C = self.codomain_grid2world
+    Pinv = self.prealign_inv
+    Dinv = self.domain_world2grid
+
+    affine_idx_in = mult_aff(Rinv, C)
+    affine_idx_out = mult_aff(Dinv, mult_aff(Pinv, C))
+    affine_disp = mult_aff(Dinv, Pinv)
+    new_backward = simplify_f(
+        self.backward,
+        affine_idx_in,
+        affine_idx_out,
+        affine_disp,
+        self.codomain_shape,
+    )
+    simplified = DiffeomorphicMap(
+        dim=self.dim,
+        disp_shape=self.disp_shape,
+        disp_grid2world=None,
+        domain_shape=self.domain_shape,
+        domain_grid2world=None,
+        codomain_shape=self.codomain_shape,
+        codomain_grid2world=None,
+        prealign=None,
+    )
+    simplified.forward = new_forward
+    simplified.backward = new_backward
+    return simplified
+
+
 def gwi_odf(gqmodel, data):
     gqi_vector = np.real(
         squared_radial_component(
@@ -243,7 +314,7 @@ def gaussian_weights(bundle, n_points=100, return_mahalnobis=False, stat=np.mean
     return weights / np.sum(weights, 0)
 
 
-def make_gif(show_m, out_path, n_frames=36, az_ang=-10):
+def make_gif(show_m, out_path, n_frames=36, az_ang=-10, duration=150):
     """
     Make a video from a Fury Show Manager.
 
@@ -263,22 +334,60 @@ def make_gif(show_m, out_path, n_frames=36, az_ang=-10):
         The angle to rotate the camera around the
         z-axis for each frame, in degrees.
         Default: -10
+
+    duration : int
+        The duration of each frame in the output GIF, in milliseconds.
+        Default: 150
     """
     video = []
 
     show_m.render()
     show_m.window.draw()
 
     with tempfile.TemporaryDirectory() as tmp_dir:
-        for ii in tqdm(range(n_frames), desc="Generating GIF"):
+        for ii in tqdm(range(n_frames), desc="Generating GIF", leave=False):
             frame_fname = f"{tmp_dir}/{ii}.png"
             show_m.screens[0].controller.rotate((radians(az_ang), 0), None)
             show_m.render()
             show_m.window.draw()
             show_m.snapshot(frame_fname)
-            video.append(frame_fname)
-
-        video = [Image.open(frame) for frame in video]
-        video[0].save(
-            out_path, save_all=True, append_images=video[1:], duration=300, loop=1
+            video.append(Image.open(frame_fname).convert("RGB"))
+
+        all_left, all_upper = float("inf"), float("inf")
+        all_right, all_lower = 0, 0
+
+        for img in video:
+            arr = np.array(img)
+            bg_color = arr[0, 0]
+
+            mask = np.any(arr != bg_color, axis=-1)
+
+            if np.any(mask):
+                rows = np.any(mask, axis=1)
+                cols = np.any(mask, axis=0)
+                ymin, ymax = np.where(rows)[0][[0, -1]]
+                xmin, xmax = np.where(cols)[0][[0, -1]]
+
+                all_left = min(all_left, xmin)
+                all_upper = min(all_upper, ymin)
+                all_right = max(all_right, xmax)
+                all_lower = max(all_lower, ymax)
+
+        if all_left < all_right:
+            crop_box = (
+                max(0, all_left),
+                max(0, all_upper),
+                min(video[0].width, all_right),
+                min(video[0].height, all_lower),
+            )
+            cropped_video = [img.crop(crop_box) for img in video]
+        else:
+            cropped_video = video
+
+        cropped_video[0].save(
+            out_path,
+            save_all=True,
+            append_images=cropped_video[1:],
+            duration=duration,
+            loop=1,
         )