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 107f56e
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 33 deletions.
75 changes: 43 additions & 32 deletions sdcflows/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@
from functools import partial
import asyncio
from pathlib import Path
from typing import Callable, List, Optional, Sequence, Tuple, Union
from typing import Callable, Sequence, Tuple

import attr
import numpy as np
Expand Down 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 @@ -481,13 +488,15 @@ def apply(
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]
-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 +516,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 107f56e

Please sign in to comment.