Commit d95b93fb authored by William Clarke's avatar William Clarke
Browse files

Add thread control to simulation script.

parent b291eb0b
......@@ -74,6 +74,10 @@ def main():
required=False, metavar='ECHOTIME',
help='Echo time value in ms for output files (no effect on simulation).')
optional.add_argument('--num_processes', type=int,
required=False, default=None,
help='Number of worker processes to use in simulation, defaults to os.cpu_count().')
# optional.add_argument('--verbose',action="store_true",
# help='spit out verbose info')
optional.add_argument('--overwrite', action="store_true",
......@@ -216,7 +220,8 @@ def main():
# Loop over the spin systems (list) in a parallel way
poolFunc = partial(runSimForMetab, seqParams=seqParams, args=args)
poolArgs = [(i, s, n) for i, (s, n) in enumerate(zip(spinsys, spinsToSim))]
pool = mp.Pool()
# handle number of processes
pool = mp.Pool(args.num_processes)
pool.starmap(poolFunc, poolArgs)
# Additional write steps for MM and LCM basis generation
......@@ -69,3 +69,27 @@ def test_sim(spinsys, seqparams, tmp_path):
assert np.allclose(basis_j[:, names_j.index('sys1')], directSim1)
assert np.allclose(basis_j[:, names_j.index('sys2')], directSim2)
def test_sim_workers(spinsys, seqparams, tmp_path):
# Test simulation number of workers
# Run the sequence on a couple of metabolites, ask for all types of files, add the reference peak
ssFile = op.join(tmp_path, 'custom_ss.json')
outfile1 = op.join(tmp_path, 'simulated_single')
outfile2 = op.join(tmp_path, 'simulated_multi')
assert not subprocess.check_call(
'-s', ssFile,
'-o', outfile1,
'-a', seqfile,
'--num_processes', '1'])
assert not subprocess.check_call(
'-s', ssFile,
'-o', outfile2,
'-a', seqfile,
'--num_processes', '2'])
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment