From c78c74e9a65cec43ea4e2822f8f7328a35928eba Mon Sep 17 00:00:00 2001
From: Oscar Esteban <code@oscaresteban.es>
Date: Wed, 16 Aug 2023 10:50:18 +0200
Subject: [PATCH] fix+sty: argument specification + run black

---
 sdcflows/transform.py                    | 83 ++++++++++++++----------
 sdcflows/workflows/apply/registration.py |  2 +-
 2 files changed, 49 insertions(+), 36 deletions(-)

diff --git a/sdcflows/transform.py b/sdcflows/transform.py
index 140c0bffb6..4181c90346 100644
--- a/sdcflows/transform.py
+++ b/sdcflows/transform.py
@@ -118,7 +118,9 @@ async def worker(
     """Create one worker and attach it to the execution loop."""
     async with semaphore:
         loop = asyncio.get_running_loop()
-        result = await loop.run_in_executor(None, func, data, coordinates, pe_info, hmc_xfm)
+        result = await loop.run_in_executor(
+            None, func, data, coordinates, pe_info, hmc_xfm
+        )
         return result
 
 
@@ -204,14 +206,16 @@ async def unwarp_parallel(
         xfm = None if xfms is None else xfms[volid]
 
         # IMPORTANT - the coordinates array must be copied every time anew per thread
-        task = asyncio.create_task(worker(
-            volume,
-            coordinates.copy(),
-            pe_info[volid],
-            xfm,
-            func,
-            semaphore,
-        ))
+        task = asyncio.create_task(
+            worker(
+                volume,
+                coordinates.copy(),
+                pe_info[volid],
+                xfm,
+                func,
+                semaphore,
+            )
+        )
         tasks.append(task)
 
     # Wait for all tasks to complete
@@ -295,15 +299,16 @@ def fit(
         finest_coeffs = listify(self.coeffs)[-1]
         approx &= not np.allclose(
             np.linalg.norm(
-                np.cross(finest_coeffs.affine[:-1, :-1].T, target_reference.affine[:-1, :-1].T),
+                np.cross(
+                    finest_coeffs.affine[:-1, :-1].T,
+                    target_reference.affine[:-1, :-1].T,
+                ),
                 axis=1,
             ),
             0,
             atol=1e-3,
         )
 
-
-
         weights = []
         coeffs = []
         if approx:
@@ -342,7 +347,9 @@ def fit(
             from nitransforms.linear import Affine
 
             _tmp_reference = nb.Nifti1Image(
-                np.zeros(target_reference.shape[:3], dtype=target_reference.get_data_dtype()),
+                np.zeros(
+                    target_reference.shape[:3], dtype=target_reference.get_data_dtype()
+                ),
                 target_reference.affine,
                 target_reference.header,
             )
@@ -364,7 +371,7 @@ def apply(
         mode: str = "constant",
         cval: float = 0.0,
         prefilter: bool = True,
-        output_dtype: str, np.dtype | None = None,
+        output_dtype: str | np.dtype | None = None,
         num_threads: int = None,
         allow_negative: bool = False,
     ):
@@ -478,16 +485,20 @@ def apply(
             axis_flip = axcodes[pe_axis] in ("LPI")
             pe_flip = pe_dir[volid].endswith("-")
 
-            pe_info.append((
-                pe_axis,
-                # Displacements are reversed if either is true (after ensuring positive cosines)
-                -ro_time[volid] if (axis_flip ^ pe_flip) else ro_time[volid]
-            ))
+            pe_info.append(
+                (
+                    pe_axis,
+                    # Displacements are reversed if either is true (after ensuring positive cosines)
+                    -ro_time[volid] if (axis_flip ^ pe_flip) else ro_time[volid],
+                )
+            )
 
         # Reference image's voxel coordinates (in voxel units)
-        voxcoords = nt.linear.Affine(
-            reference=moving
-        ).reference.ndindex.reshape((ndim, *data.shape[:ndim])).astype("float32")
+        voxcoords = (
+            nt.linear.Affine(reference=moving)
+            .reference.ndindex.reshape((ndim, *data.shape[:ndim]))
+            .astype("float32")
+        )
 
         # Convert head-motion transforms to voxel-to-voxel:
         if xfms is not None:
@@ -507,19 +518,21 @@ def apply(
             )
 
         # Resample
-        resampled = asyncio.run(unwarp_parallel(
-            data,
-            voxcoords,
-            self.mapped.get_fdata(dtype="float32"),  # fieldmap in Hz
-            pe_info,
-            xfms,
-            output_dtype=output_dtype,
-            order=order,
-            mode=mode,
-            cval=cval,
-            prefilter=prefilter,
-            max_concurrent=num_threads or min(os.cpu_count(), 12),
-        ))
+        resampled = asyncio.run(
+            unwarp_parallel(
+                data,
+                voxcoords,
+                self.mapped.get_fdata(dtype="float32"),  # fieldmap in Hz
+                pe_info,
+                xfms,
+                output_dtype=output_dtype,
+                order=order,
+                mode=mode,
+                cval=cval,
+                prefilter=prefilter,
+                max_concurrent=num_threads or min(os.cpu_count(), 12),
+            )
+        )
 
         if not allow_negative:
             resampled[resampled < 0] = cval
diff --git a/sdcflows/workflows/apply/registration.py b/sdcflows/workflows/apply/registration.py
index 7bb8ee4fcf..a82f8c23a5 100644
--- a/sdcflows/workflows/apply/registration.py
+++ b/sdcflows/workflows/apply/registration.py
@@ -147,7 +147,7 @@ def init_coeff2epi_wf(
 
     if write_coeff:
         warn(
-            "SDCFlows does not tinker with the coefficients file anymore, " 
+            "SDCFlows does not tinker with the coefficients file anymore, "
             "the `write_coeff` parameter will be removed in a future release."
         )