diff --git a/trackpy/linking/linking.py b/trackpy/linking/linking.py index 5b9e111e..da67f4e6 100644 --- a/trackpy/linking/linking.py +++ b/trackpy/linking/linking.py @@ -5,6 +5,8 @@ import warnings import logging import itertools, functools +from concurrent import futures +from multiprocessing import cpu_count import numpy as np @@ -499,12 +501,19 @@ def next_level(self, coords, t, extra_data=None): def assign_links(self): spl, dpl = [], [] - for source_set, dest_set in self.subnets: + + def link_subnet(subnet_sets): + source_set, dest_set = subnet_sets for sp in source_set: sp.forward_cands.sort(key=lambda x: x[1]) - sn_spl, sn_dpl = self.subnet_linker(source_set, dest_set, - self.search_range) + return self.subnet_linker(source_set, dest_set, + self.search_range) + + workers = futures.ThreadPoolExecutor(max_workers=cpu_count()) + futs = [workers.submit(link_subnet, sub) for sub in self.subnets] + for fut in futures.as_completed(futs): + sn_spl, sn_dpl = fut.result() spl.extend(sn_spl) dpl.extend(sn_dpl) diff --git a/trackpy/linking/subnetlinker.py b/trackpy/linking/subnetlinker.py index 1753e2f1..ce7d9f56 100644 --- a/trackpy/linking/subnetlinker.py +++ b/trackpy/linking/subnetlinker.py @@ -238,7 +238,7 @@ def numba_link(s_sn, dest_size, search_range, max_size=30, diag=False): dest_results = [dcands[i] if i >= 0 else None for i in best_assignments] return source_results, dest_results -@try_numba_jit(nopython=True) +@try_numba_jit(nopython=True, nogil=True) def _numba_subnet_norecur(ncands, candsarray, dists2array, cur_assignments, cur_sums, tmp_assignments, best_assignments): """Find the optimal track assignments for a subnetwork, without recursion.