def __generate_charges(
self,
graphs: List[Tuple[int, nx.Graph]],
color_key: str,
traceable: bool = False,
versioing: bool = False) -> Dict[int, Dict[str, List[float]]]:
"""Generate charges for all shell sizes and neighborhoods."""
if not versioing:
charges = defaultdict(lambda: defaultdict(list))
else:
charges = defaultdict(lambda: defaultdict(_VersioningList))
if traceable:
Worker = _TraceableChargeWorker
else:
Worker = _ChargeWorker
for shell in range(self.__min_shell, self.__max_shell + 1):
with MultiProcessor(Worker, (shell, color_key)) as mp:
for c in mp.processed(graphs, 'shell %d' % shell):
for key, values in c.items():
charges[shell][key] += values
for shell in range(self.__min_shell, self.__max_shell + 1):
for key, values in charges[shell].items():
charges[shell][key].sort()
return charges
def __make_canons(self, graphs: List[Tuple[int, nx.Graph]],
color_key: str) -> Dict[int, str]:
"""Canonicalize the given graphs using Nauty."""
canons = dict()
with MultiProcessor(_CanonicalizationWorker, color_key) as mp:
for molid, canon in mp.processed(graphs):
canons[molid] = canon
return canons
def __read_graphs(self, molids: List[int], data_location: str, ext: str,
data_type: IOType) -> List[Tuple[int, nx.Graph]]:
"""Read graphs from a directory of input files."""
graphs = []
with MultiProcessor(_ReadWorker,
(data_location, ext, data_type)) as mp:
for molid, graph in mp.processed(molids, 'reading files'):
graphs.append((molid, graph))
return graphs
def test_processing_parallel():
mp = MultiProcessor(ProcClassAdder, 5, num_processes=2)
result = list(mp.processed(range(100)))
mp.shutdown()
assert set(result) == set(range(5, 105))
def test_scope_guard():
with MultiProcessor(ProcClass, num_processes=2) as mp:
assert mp is not None
def test_processing_serial():
mp = MultiProcessor(ProcClassAdder, 5, num_processes=1)
result = list(mp.processed(range(100)))
mp.shutdown()
assert result == list(range(5, 105))
def test_create_error_handling():
with pytest.raises(RuntimeError):
mp = MultiProcessor(ProcClassInitTester, (1, 2), 1)
def test_create_initargs3():
mp = MultiProcessor(ProcClassInitTester, (1, '3'), 1)
mp.shutdown()
def test_create_initargs2():
mp = MultiProcessor(ProcClassInitTester, 1, 1)
mp.shutdown()
def test_create_initargs0():
mp = MultiProcessor(ProcClassInitZeroTester, 0, 1)
mp.shutdown()
def test_create2():
mp = MultiProcessor(ProcClass, num_processes=2)
mp.shutdown()