Source code for xicsrt.xicsrt_multiprocessing

# -*- coding: utf-8 -*-
"""
.. Authors
    Novimir pablant <npablant@pppl.gov>
"""

from multiprocessing import Pool
from xicsrt.xicsrt_raytrace import *

m_log = mirlogging.getLogger(__name__)

[docs] def raytrace(config, processes=None): """ Perform a series of ray tracing runs using the `multiprocessing` cpu pool. Each run will rebuild all objects, reset the random seed and then perform the requested number of iterations. If the option 'save_images' is set, then images will be saved at the completion of each run. Also see :func:`~xicsrt.xicsrt_raytrace.raytrace` for a single process version of this routine. """ profiler.start('raytrace_multiprocessing') # Update the default config with the user config. config = xicsrt_config.get_config(config) check_config(config) num_runs = config['general']['number_of_runs'] random_seed = config['general']['random_seed'] mp_result_list = [] output_list = [] with Pool(processes) as pool: # loop through each configuration in the configuration input file # and add a new run into the pool. for ii in range(num_runs): m_log.info('Adding run to pool: {} of {}'.format(ii + 1, num_runs)) # Make a copy of the configuration. config_run = deepcopy(config) config_run['general']['output_run_suffix'] = '{:04d}'.format(ii) # Make sure each run uses a unique random seed. if random_seed is not None: random_seed += ii config_run['general']['random_seed'] = random_seed arg = (config_run, True) mp_result = pool.apply_async(raytrace_single, arg) mp_result_list.append(mp_result) pool.close() pool.join() profiler.start('mp: gathering') # Gather all the results together. for mp_result in mp_result_list: output = mp_result.get() output_list.append(output) profiler.stop('mp: gathering') output = combine_raytrace(output_list) # Reset the configuration options that were unique to the individual runs. output['config']['general']['output_run_suffix'] = config['general']['output_run_suffix'] output['config']['general']['random_seed'] = config['general']['output_run_suffix'] if config['general']['save_config']: xicsrt_io.save_config(output['config']) if config['general']['save_images']: xicsrt_io.save_images(output) if config['general']['save_results']: xicsrt_io.save_results(output) if config['general']['print_results']: print_raytrace(output) profiler.stop('raytrace_multiprocessing') return output