Source code for xicsrt.xicsrt_raytrace

# -*- coding: utf-8 -*-
"""
.. Authors:
    Novimir Pablant <npablant@pppl.gov>
    Yevgeniy Yakusevich <eugenethree@gmail.com>
    James Kring <jdk0026@tigermail.auburn.edu>

Entry point to XICSRT.
Contains the main functions that are called to perform raytracing.
"""

import numpy as np

import os

from copy import deepcopy

from xicsrt.util import mirlogging
from xicsrt.util import profiler

from xicsrt import xicsrt_config
from xicsrt import xicsrt_io
from xicsrt.objects._Dispatcher import Dispatcher
from xicsrt.objects._RayArray import RayArray

m_log = mirlogging.getLogger(__name__)

[docs] def raytrace(config): """ Perform a series of ray tracing runs. 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. The saving of these run images is one reason to use this routine rather than just increasing the number of iterations: periodic outputs during long computations. Also see :func:`~xicsrt.xicsrt_multiprocessing.raytrace` for a multiprocessing version of this routine. """ profiler.start('raytrace') # Update the default config with the user config. config = xicsrt_config.get_config(config) check_config(config) # Make local copies of some options. num_runs = config['general']['number_of_runs'] random_seed = config['general']['random_seed'] output_list = [] for ii in range(num_runs): m_log.info('Starting run: {} of {}'.format(ii + 1, num_runs)) 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 iteration = raytrace_single(config_run, _internal=True) output_list.append(iteration) 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']['random_seed'] 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') return output
[docs] def raytrace_single(config, _internal=False): """ Perform a single raytrace run consisting of multiple iterations. If history is enabled, sort the rays into those that are detected and those that are lost (found and lost). The found ray history will be returned in full. The lost ray history will be truncated to allow analysis of lost ray pattern while still limiting memory usage. private keywords ================ _internal : bool (False) Used when calling this function from `raytrace` as part of the execution of multiple runs. Controls how `history_max_lost` is handled along with how `save_config` and `save_results` are interpreted. """ profiler.start('raytrace_single') # Update the default config with the user config. config = xicsrt_config.config_to_numpy(config) config = xicsrt_config.get_config(config) check_config(config) m_log.info('Seeding np.random with {}'.format(config['general']['random_seed'])) np.random.seed(config['general']['random_seed']) num_iter = config['general']['number_of_iter'] max_lost_iter = int(config['general']['history_max_lost']/num_iter) if _internal: max_lost_iter = max_lost_iter//config['general']['number_of_runs'] # Save at least one lost ray even if this exceeds history_max_lost. max_lost_iter = max(int(max_lost_iter), 1) # Setup the dispatchers. if 'filters' in config: m_log.debug("Creating filters") filters = Dispatcher(config, 'filters') filters.instantiate() filters.setup() filters.initialize() config['filters'] = filters.get_config() else: filters = None m_log.debug("Creating sources") sources = Dispatcher(config, 'sources') sources.instantiate() sources.apply_filters(filters) sources.setup() sources.check_param() sources.initialize() config['sources'] = sources.get_config() m_log.debug("Creating optics") optics = Dispatcher(config, 'optics') optics.instantiate() optics.apply_filters(filters) optics.setup() optics.check_param() optics.initialize() config['optics'] = optics.get_config() # Do the actual raytracing output_list = [] for ii in range(num_iter): m_log.info('Starting iteration: {} of {}'.format(ii + 1, num_iter)) single = _raytrace_iter(config, sources, optics) sorted = _sort_raytrace(single, max_lost=max_lost_iter) output_list.append(sorted) output = combine_raytrace(output_list) if _internal is False: if config['general']['print_results']: print_raytrace(output) if config['general']['save_config']: xicsrt_io.save_config(output['config']) if config['general']['save_results']: xicsrt_io.save_results(output) if config['general']['save_images']: xicsrt_io.save_images(output) profiler.stop('raytrace_single') # profiler.report() return output
[docs] def _raytrace_iter(config, sources, optics): """ Perform a single iteration of raytracing with the given sources and optics. The returned rays are unsorted. """ profiler.start('_raytrace_iter') # Setup local names for a few config entries. # This is only to make the code below more readable. keep_meta = config['general']['keep_meta'] keep_images = config['general']['keep_images'] keep_history = config['general']['keep_history'] m_log.debug('Generating rays') rays = sources.generate_rays(keep_history=keep_history) m_log.debug('Raytracing optics') rays = optics.trace(rays, keep_history=keep_history, keep_images=keep_images) # Combine sources and optics outputs. meta = dict() image = dict() history = dict() if keep_meta: for key in sources.meta: meta[key] = sources.meta[key] for key in optics.meta: meta[key] = optics.meta[key] if keep_images: for key in sources.image: image[key] = sources.image[key] for key in optics.image: image[key] = optics.image[key] if keep_history: for key in sources.history: history[key] = sources.history[key] for key in optics.history: history[key] = optics.history[key] output = dict() output['config'] = config output['meta'] = meta output['image'] = image output['history'] = history profiler.stop('_raytrace_iter') return output
[docs] def _sort_raytrace(input, max_lost=None): """ Sort the rays into 'lost' and 'found' rays, then truncate the number of lost rays. """ if max_lost is None: max_lost = 1000 profiler.start('_sort_raytrace') output = dict() output['config'] = input['config'] output['total'] = dict() output['total']['meta'] = dict() output['total']['image'] = dict() output['found'] = dict() output['found']['meta'] = dict() output['found']['history'] = dict() output['lost'] = dict() output['lost']['meta'] = dict() output['lost']['history'] = dict() output['total']['meta'] = input['meta'] output['total']['image'] = input['image'] if len(input['history']) > 0: key_opt_list = list(input['history'].keys()) key_opt_last = key_opt_list[-1] w_found = np.flatnonzero(input['history'][key_opt_last]['mask']) w_lost = np.flatnonzero(np.invert(input['history'][key_opt_last]['mask'])) # Save only a portion of the lost rays so that our lost history does # not become too large. max_lost = min(max_lost, len(w_lost)) index_lost = np.arange(len(w_lost)) np.random.shuffle(index_lost) w_lost = w_lost[index_lost[:max_lost]] for key_opt in key_opt_list: output['found']['history'][key_opt] = dict() output['lost']['history'][key_opt] = dict() for key_ray in input['history'][key_opt]: output['found']['history'][key_opt][key_ray] = input['history'][key_opt][key_ray][w_found] output['lost']['history'][key_opt][key_ray] = input['history'][key_opt][key_ray][w_lost] profiler.stop('_sort_raytrace') return output
[docs] def combine_raytrace(input_list, keep_images=True, components=None): """ Produce a combined results dictionary from a list of raytrace results. Keywords -------- keep_images: bool (True) Control whether to combine images. Useful when combining outputs which use different optics sizes. components: list (None) A list of specific components to combine. Useful when not all components are needed in the final output. Example ------- results_1 = xicsrt.raytrace(config_1) results_2 = xicsrt.raytrace(config_2) results = xicsrt_raytrace.combine_raytrace([results_1, results_2]) """ profiler.start('combine_raytrace') output = dict() output['config'] = input_list[0]['config'] output['total'] = dict() output['total']['meta'] = dict() output['total']['image'] = dict() output['found'] = dict() output['found']['meta'] = dict() output['found']['history'] = dict() output['lost'] = dict() output['lost']['meta'] = dict() output['lost']['history'] = dict() num_iter = len(input_list) if components is None: key_opt_list = list(input_list[0]['total']['meta'].keys()) else: key_opt_list = components key_opt_last = key_opt_list[-1] # Combine the meta data. for key_opt in key_opt_list: output['total']['meta'][key_opt] = dict() key_meta_list = list(input_list[0]['total']['meta'][key_opt].keys()) for key_meta in key_meta_list: output['total']['meta'][key_opt][key_meta] = 0 for ii_iter in range(num_iter): output['total']['meta'][key_opt][key_meta] += input_list[ii_iter]['total']['meta'][key_opt][key_meta] # Combine the images. if keep_images: for key_opt in key_opt_list: if key_opt in input_list[0]['total']['image']: if input_list[0]['total']['image'][key_opt] is not None: # Check shape compatibility shape = input_list[0]['total']['image'][key_opt].shape shape_match = True for ii_iter in range(num_iter): if not input_list[ii_iter]['total']['image'][key_opt].shape == shape: shape_match = False if shape_match: # Shapes match, combine images output['total']['image'][key_opt] = np.zeros(input_list[0]['total']['image'][key_opt].shape) for ii_iter in range(num_iter): output['total']['image'][key_opt] += input_list[ii_iter]['total']['image'][key_opt] else: m_log.warning('Image dimensions do not match. Cannot combine images.') output['total']['image'][key_opt] = None else: output['total']['image'][key_opt] = None # Combine all the histories. if len(input_list[0]['found']['history']) > 0: final_num_found = 0 final_num_lost = 0 for ii_run in range(num_iter): final_num_found += len(input_list[ii_run]['found']['history'][key_opt_last]['mask']) final_num_lost += len(input_list[ii_run]['lost']['history'][key_opt_last]['mask']) rays_found_temp = RayArray() rays_found_temp.zeros(final_num_found) rays_lost_temp = RayArray() rays_lost_temp.zeros(final_num_lost) for key_opt in key_opt_list: output['found']['history'][key_opt] = rays_found_temp.copy() output['lost']['history'][key_opt] = rays_lost_temp.copy() index_found = 0 index_lost = 0 for ii_run in range(num_iter): num_found = len(input_list[ii_run]['found']['history'][key_opt_last]['mask']) num_lost = len(input_list[ii_run]['lost']['history'][key_opt_last]['mask']) for key_opt in key_opt_list: for key_ray in output['found']['history'][key_opt]: output['found']['history'][key_opt][key_ray][index_found:index_found + num_found] = ( input_list[ii_run]['found']['history'][key_opt][key_ray][:]) output['lost']['history'][key_opt][key_ray][index_lost:index_lost + num_lost] = ( input_list[ii_run]['lost']['history'][key_opt][key_ray][:]) index_found += num_found index_lost += num_lost profiler.stop('combine_raytrace') return output
[docs] def check_config(config): """ Check the general section of the configuration dictionary. """ # Check if anything needs to be saved. do_save = False for key in config['general']: if 'save' in key: if config['general'][key]: do_save = True if do_save: if not xicsrt_io.path_exists(config['general']['output_path']): if not config['general']['make_directories']: raise Exception('Output directory does not exist. Create directory or set make_directories to True.')