Skip to content

Commit

Permalink
fix+sty: argument specification + run black
Browse files Browse the repository at this point in the history
  • Loading branch information
oesteban committed Aug 16, 2023
1 parent 2e72608 commit c78c74e
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 36 deletions.
83 changes: 48 additions & 35 deletions sdcflows/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
)
Expand All @@ -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,
):
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion sdcflows/workflows/apply/registration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)

Expand Down

0 comments on commit c78c74e

Please sign in to comment.