Skip to content

Commit 1ca7763

Browse files
committed
Switch to Processes for voronoi cell volume calculation
Ref: Speed up dcf calculation #169 Scipy.spatial Qhull does not seem to work in parallel in multiple threads, using processes makes it~4x faster on my notebook. Would still prefer a gpu version.
1 parent 025bc79 commit 1ca7763

File tree

1 file changed

+8
-3
lines changed

1 file changed

+8
-3
lines changed

src/mrpro/data/_DcfData.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from __future__ import annotations
1616

1717
import dataclasses
18-
from concurrent.futures import ThreadPoolExecutor
18+
from concurrent.futures import ProcessPoolExecutor
1919
from functools import reduce
2020
from itertools import product
2121

@@ -30,6 +30,10 @@
3030
UNIQUE_ROUNDING_DECIMALS = 15
3131

3232

33+
def _volume(v):
34+
return ConvexHull(v).volume
35+
36+
3337
@dataclasses.dataclass(slots=True, frozen=False)
3438
class DcfData:
3539
"""Density compensation data (DcfData) class."""
@@ -114,8 +118,9 @@ def _dcf_2d3d_voronoi(traj: torch.Tensor) -> torch.Tensor:
114118
regions = [vdiagram.regions[r] for r in vdiagram.point_region[: -len(corner_points)]] # Ignore corner points
115119
vertices = [vdiagram.vertices[region] for region in regions]
116120

117-
# Calculate volume/area of voronoi cells using threads (ConvexHull is thread safe and drops the GIL)
118-
future = ThreadPoolExecutor(max_workers=torch.get_num_threads()).map(lambda v: ConvexHull(v).volume, vertices)
121+
# Calculate volume/area of voronoi cells using processes, as this is a very time-consuming operation
122+
# and ConvexHull is singlethreaded and does not seem to drop the GIL
123+
future = ProcessPoolExecutor(max_workers=torch.get_num_threads()).map(_volume, vertices, chunksize=100)
119124
dcf = np.array(list(future))
120125

121126
# Get outliers (i.e. voronoi cell which are unbound) and set them to a reasonable value

0 commit comments

Comments
 (0)