diff --git a/mcdc/code_factory/gpu/adapt.py b/mcdc/code_factory/gpu/adapt.py deleted file mode 100644 index 4d504f2e..00000000 --- a/mcdc/code_factory/gpu/adapt.py +++ /dev/null @@ -1,613 +0,0 @@ -import importlib -import inspect -import numba -import numpy as np - -from numba import njit, jit, types -from numba.extending import intrinsic - -if importlib.util.find_spec("harmonize") is None: - HAS_HARMONIZE = False -else: - import harmonize as harm - - HAS_HARMONIZE = True - -#### - - -import mcdc.config as config - -from mcdc.print_ import print_error - -# ============================================================================= -# Error Messangers -# ============================================================================= - - -def unknown_target(target): - print_error(f"ERROR: Unrecognized target '{target}'") - - -# ============================================================================= -# uintp/voidptr casters -# ============================================================================= - - -@intrinsic -def cast_uintp_to_voidptr(typingctx, src): - # check for accepted types - if isinstance(src, types.Integer): - # create the expected type signature - result_type = types.voidptr - sig = result_type(types.uintp) - - # defines the custom code generation - def codegen(context, builder, signature, args): - # llvm IRBuilder code here - [src] = args - rtype = signature.return_type - llrtype = context.get_value_type(rtype) - return builder.inttoptr(src, llrtype) - - return sig, codegen - - -@intrinsic -def cast_voidptr_to_uintp(typingctx, src): - # check for accepted types - if isinstance(src, types.RawPointer): - # create the expected type signature - result_type = types.uintp - sig = result_type(types.voidptr) - - # defines the custom code generation - def codegen(context, builder, signature, args): - # llvm IRBuilder code here - [src] = args - rtype = signature.return_type - llrtype = context.get_value_type(rtype) - return builder.ptrtoint(src, llrtype) - - return sig, codegen - - -@njit() -def uintp_to_voidptr(value): - val = numba.uintp(value) - return cast_uintp_to_voidptr(val) - - -@njit() -def voidptr_to_uintp(value): - return cast_voidptr_to_uintp(value) - - -def leak(arg): - pass - - -@intrinsic -def leak_inner(typingctx, kind): - def codegen(context, builder, signature, args): - context.nrt.incref(builder, kind, args[0]) - - return numba.void(kind), codegen - - -@numba.extending.overload(leak) -def leak_overload(arg): - def impl(arg): - leak_inner(arg) - - return impl - - -# ============================================================================= -# Generic GPU/CPU Local Array Variable Constructors -# ============================================================================= - - -def local_array(shape, dtype): - return np.zeros(shape, dtype=dtype) - - -@numba.extending.type_callable(local_array) -def type_local_array(context): - - from numba.core.typing.npydecl import parse_dtype, parse_shape - - if isinstance(context, numba.core.typing.context.Context): - - # Function repurposed from Numba's ol_np_empty. - def typer(shape, dtype): - numba.np.arrayobj._check_const_str_dtype("empty", dtype) - - # Only integer literals and tuples of integer literals are valid - # shapes - if isinstance(shape, types.Integer): - if not isinstance(shape, types.IntegerLiteral): - raise numba.core.errors.UnsupportedError( - f"Integer shape type {shape} is not literal." - ) - elif isinstance(shape, (types.Tuple, types.UniTuple)): - if any([not isinstance(s, types.IntegerLiteral) for s in shape]): - raise numba.core.errors.UnsupportedError( - f"At least one element of shape tuple type{shape} is not an integer literal." - ) - else: - raise numba.core.errors.UnsupportedError( - f"Shape is of unsupported type {shape}." - ) - - # No default arguments. - nb_dtype = parse_dtype(dtype) - nb_shape = parse_shape(shape) - - if nb_dtype is not None and nb_shape is not None: - retty = types.Array(dtype=nb_dtype, ndim=nb_shape, layout="C") - # Inlining the signature construction from numpy_empty_nd - sig = retty(shape, dtype) - return sig - else: - msg = f"Cannot parse input types to function np.empty({shape}, {dtype})" - raise numba.errors.TypingError(msg) - - return typer - - elif isinstance(context, numba.cuda.target.CUDATypingContext): - - # Function repurposed from Numba's Cuda_array_decl. - def typer(shape, dtype): - - # Only integer literals and tuples of integer literals are valid - # shapes - if isinstance(shape, types.Integer): - if not isinstance(shape, types.IntegerLiteral): - return None - elif isinstance(shape, (types.Tuple, types.UniTuple)): - if any([not isinstance(s, types.IntegerLiteral) for s in shape]): - return None - else: - return None - - ndim = parse_shape(shape) - nb_dtype = parse_dtype(dtype) - if nb_dtype is not None and ndim is not None: - return types.Array(dtype=nb_dtype, ndim=ndim, layout="C") - - return typer - - elif isinstance(context, numba.hip.target.HIPTypingContext): - - def typer(shape, dtype): - # Only integer literals and tuples of integer literals are valid - # shapes - if isinstance(shape, types.Integer): - if not isinstance(shape, types.IntegerLiteral): - return None - elif isinstance(shape, (types.Tuple, types.UniTuple)): - if any([not isinstance(s, types.IntegerLiteral) for s in shape]): - return None - else: - return None - - ndim = parse_shape(shape) - nb_dtype = parse_dtype(dtype) - if nb_dtype is not None and ndim is not None: - result = types.Array(dtype=nb_dtype, ndim=ndim, layout="C") - return result - - return typer - - else: - raise numba.core.errors.UnsupportedError( - f"Unsupported target context {context}." - ) - - -@numba.extending.lower_builtin(local_array, types.IntegerLiteral, types.Any) -def builtin_local_array(context, builder, sig, args): - - shape, dtype = sig.args - - from numba.core.typing.npydecl import parse_dtype, parse_shape - import numba.np.arrayobj as arrayobj - - if isinstance(context, numba.core.cpu.CPUContext): - - # No default arguments. - nb_dtype = parse_dtype(dtype) - nb_shape = parse_shape(shape) - - retty = types.Array(dtype=nb_dtype, ndim=nb_shape, layout="C") - - # In ol_np_empty, the reference type of the array is fed into the - # signatrue as a third argument. This third argument is not used by - # _parse_empty_args. - sig = retty(shape, dtype) - - arrtype, shapes = arrayobj._parse_empty_args(context, builder, sig, args) - ary = arrayobj._empty_nd_impl(context, builder, arrtype, shapes) - - return ary._getvalue() - elif isinstance(context, numba.cuda.target.CUDATargetContext): - length = sig.args[0].literal_value - dtype = parse_dtype(sig.args[1]) - return numba.cuda.cudaimpl._generic_array( - context, - builder, - shape=(length,), - dtype=dtype, - symbol_name="_cudapy_harm_lmem", - addrspace=numba.cuda.cudadrv.nvvm.ADDRSPACE_LOCAL, - can_dynsized=False, - ) - elif isinstance(context, numba.hip.target.HIPTargetContext): - length = sig.args[0].literal_value - dtype = parse_dtype(sig.args[1]) - result = numba.hip.typing_lowering.hip.lowering._generic_array( - context, - builder, - shape=(length,), - dtype=dtype, - symbol_name="_HIPpy_lmem", - addrspace=numba.hip.amdgcn.ADDRSPACE_LOCAL, - can_dynsized=False, - ) - return result - else: - raise numba.core.errors.UnsupportedError( - f"Unsupported target context {context}." - ) - - -# ============================================================================= -# Decorators -# ============================================================================= - -target_rosters = {} - -late_jit_roster = set() - -do_nothing_id = 0 - - -do_nothing_id = 0 - - -def generate_do_nothing(arg_count, crash_on_call=None): - """ - Create a no-op function (or one that always asserts) that takes - `arg_count` positional arguments. - - On Python 3.13+, we can't rely on exec() mutating the current frame's - locals in a way eval() can see, so we execute into an explicit - namespace and pull the function out of that. - """ - global do_nothing_id - - name = f"do_nothing_{do_nothing_id}" - args = ", ".join([f"arg_{i}" for i in range(arg_count)]) - source = f"def {name}({args}):\n" - if crash_on_call is not None: - source += f" assert False, '{crash_on_call}'\n" - else: - source += " pass\n" - - ns = {} - exec(source, globals(), ns) - result = ns[name] - - do_nothing_id += 1 - return result - - -def overwrite_func(func, revised_func): - mod_name = func.__module__ - fn_name = func.__name__ - new_fn_name = revised_func.__name__ - module = __import__(mod_name, fromlist=[fn_name]) - setattr(module, fn_name, revised_func) - - -blankout_roster = {} - - -def blankout_fn(func): - global blankout_roster - - mod_name = func.__module__ - fn_name = func.__name__ - id = (mod_name, fn_name) - - if id not in blankout_roster: - global do_nothing_id - name = func.__name__ - arg_count = len(inspect.signature(func).parameters) - blankout_roster[id] = generate_do_nothing( - arg_count, crash_on_call=f"blankout fn for {name} should never be called" - ) - - blank = blankout_roster[id] - - return blank - - -def for_(target, on_target=[]): - def for_inner(func): - global target_rosters - mod_name = func.__module__ - fn_name = func.__name__ - params = inspect.signature(func).parameters - if target not in target_rosters: - target_rosters[target] = {} - target_rosters[target][(mod_name, fn_name)] = func - - param_str = ", ".join(p for p in params) - jit_str = ( - f"def jit_func({param_str}):\n" - f" global target_rosters\n" - f" return target_rosters['{target}'][('{mod_name}','{fn_name}')]" - ) - - # Execute into an explicit namespace so we can reliably retrieve jit_func - ns = {} - exec(jit_str, globals(), ns) - result = ns["jit_func"] - - blank = blankout_fn(func) - if target == "gpu": - numba.core.extending.overload(blank, target=target)(result) - else: - numba.core.extending.overload(blank, target=target)(result) - return blank - - return for_inner - - -def for_cpu(on_target=[]): - return for_("cpu", on_target=on_target) - - -def for_gpu(on_target=[]): - return for_("gpu", on_target=on_target) - - -def jit_on_target(): - def jit_on_target_inner(func): - late_jit_roster.add(func) - return func - - return jit_on_target_inner - - -def nopython_mode(is_on): - if is_on: - return - if not isinstance(target_rosters["cpu"], dict): - return - - for impl in target_rosters["cpu"].values(): - overwrite_func(impl, impl) - - -# ============================================================================= -# GPU Type / Extern Functions Forward Declarations -# ============================================================================= - - -SIMPLE_ASYNC = True - -none_type = None -mcdc_global_type = None -mcdc_data_type = None -mcdc_shared_type = None -state_spec = None -mcdc_global_gpu = None -mcdc_data_gpu = None -group_gpu = None -thread_gpu = None -particle_gpu = None -prep_gpu = None -step_async = None -halt_early = None -find_cell_async = None -tally_width = None -tally_length = None -tally_size = None -tally_shape_literal = None - - -def gpu_forward_declare( - args, data_shape, global_type, particle_type, particle_data_type -): - - if args.gpu_rocm_path != None: - harm.config.set_rocm_path(args.gpu_rocm_path) - - if args.gpu_cuda_path != None: - harm.config.set_cuda_path(args.gpu_cuda_path) - - global none_type, mcdc_global_type, mcdc_data_type, mcdc_shared_type - global state_spec - global mcdc_global_gpu, mcdc_data_gpu - global group_gpu, thread_gpu - global particle_gpu, particle_record_gpu - global step_async, find_cell_async, halt_early - global tally_width, tally_length, tally_size - - tally_size = data_shape[0] * 8 - - global tally_shape_literal - tally_shape_literal = data_shape - - none_type = numba.from_dtype(np.dtype([])) - mcdc_global_type = numba.types.Array(numba.from_dtype(global_type), (1,), "C") - # mcdc_global_type = numba.from_dtype(global_type) - - tally_dims = len(data_shape) - mcdc_data_type = numba.types.Array(numba.float64, tally_dims, "C") - state_spec = ( - { - "global": mcdc_global_type, - "data": mcdc_data_type, - }, - none_type, - none_type, - ) - access_fns = harm.RuntimeSpec.access_fns(state_spec) - mcdc_global_gpu = access_fns["device"]["global"]["indirect"] - mcdc_data_gpu = access_fns["device"]["data"]["direct"] - group_gpu = access_fns["group"] - thread_gpu = access_fns["thread"] - particle_gpu = numba.from_dtype(particle_type) - particle_record_gpu = numba.from_dtype(particle_data_type) - - def step(prog: numba.uintp, P: particle_gpu): - pass - - def find_cell(prog: numba.uintp, P: particle_gpu): - pass - - import harmonize - - step_async, find_cell_async = harmonize.RuntimeSpec.async_dispatch(step, find_cell) - interface = harmonize.RuntimeSpec.program_interface() - halt_early = interface["halt_early"] - - -# ============================================================================= -# Seperate GPU/CPU Functions to Target Different Platforms -# ============================================================================= - - -def mcdc_global(prog): - return prog - - -@for_cpu() -def mcdc_global(prog): - return prog - - -@for_gpu() -def mcdc_global(prog): - return mcdc_global_gpu(prog) - - -@for_cpu() -def mcdc_data(prog): - return None - - -@for_gpu() -def mcdc_data(prog): - return mcdc_data_gpu(prog) - - -@for_cpu() -def group(prog): - return prog - - -@for_gpu() -def group(prog): - return group_gpu(prog) - - -@for_cpu() -def thread(prog): - return prog - - -@for_gpu() -def thread(prog): - return thread_gpu(prog) - - -@for_cpu() -def global_add(ary, idx, val): - result = ary[idx] - ary[idx] += val - return result - - -@for_gpu() -def global_add(ary, idx, val): - return harm.array_atomic_add(ary, idx, val) - - -@for_cpu() -def global_max(ary, idx, val): - result = ary[idx] - if ary[idx] < val: - ary[idx] = val - return result - - -@for_gpu() -def global_max(ary, idx, val): - return harm.array_atomic_max(ary, idx, val) - - -# ========================================================================= -# Program Specifications -# ========================================================================= - -state_spec = None -one_event_fns = None -multi_event_fns = None - - -device_gpu, group_gpu, thread_gpu = None, None, None -iterate_async = None - - -def make_spec(target): - global state_spec, one_event_fns, multi_event_fns - global device_gpu, group_gpu, thread_gpu - global iterate_async - if target == "gpu": - state_spec = (dev_state_type, grp_state_type, thd_state_type) - one_event_fns = [iterate] - # multi_event_fns = [source,move,scattering,fission,leakage,bcollision] - device_gpu, group_gpu, thread_gpu = harm.RuntimeSpec.access_fns(state_spec) - (iterate_async,) = harm.RuntimeSpec.async_dispatch(iterate) - elif target != "cpu": - unknown_target(target) - - -@njit -def empty_base_func(prog): - pass - - -def make_gpu_loop( - state_spec, - work_make_fn, - step_fn, - check_fn, - arg_type, - initial_fn=empty_base_func, - final_fn=empty_base_func, -): - async_fn_list = [step_fn] - device_gpu, group_gpu, thread_gpu = harm.RuntimeSpec.access_fns(state_spec) - - def make_work(prog: numba.uintp) -> numba.boolean: - return work_make_fn(prog) - - def initialize(prog: numba.uintp): - initial_fn(prog) - - def finalize(prog: numba.uintp): - final_fn(prog) - - def step(prog: numba.uintp, arg: arg_type): - - step_async() - - (step_async,) = harm.RuntimeSpec.async_dispatch(step) - - pass diff --git a/mcdc/code_factory/gpu/program_builder.py b/mcdc/code_factory/gpu/program_builder.py index d6d0e4ff..96170404 100644 --- a/mcdc/code_factory/gpu/program_builder.py +++ b/mcdc/code_factory/gpu/program_builder.py @@ -1,4 +1,3 @@ -import harmonize import numba as nb import numba.extending as nbxt import numpy as np @@ -8,19 +7,80 @@ #### import mcdc.config as config -import mcdc.numba_types as type_ # ====================================================================================== -# Build GPU program +# Transport function adapter # ====================================================================================== -# For teardown -src_free_program = lambda pointer: None -free_state = lambda pointer: None +def adapt_transport_functions(): + global access_simulation -def build_gpu_program(mcdc_container, data): - global src_free_program, free_state + import mcdc.code_factory.gpu.transport as gpu_transport + import mcdc.transport as transport + + transport.util.access_simulation = access_simulation + + # TODO: Make the following automatic + transport.geometry.interface.report_lost_particle = ( + gpu_transport.geometry.interface.report_lost_particle + ) + transport.particle_bank.bank_active_particle = ( + gpu_transport.particle_bank.bank_active_particle + ) + transport.particle_bank.report_full_bank = ( + gpu_transport.particle_bank.report_full_bank + ) + transport.particle_bank.report_empty_bank = ( + gpu_transport.particle_bank.report_empty_bank + ) + transport.util.atomic_add = gpu_transport.util.atomic_add + transport.util.local_array = gpu_transport.util.local_array + + +def adapt_transport_functions_post_setup(): + import mcdc.code_factory.gpu.transport as gpu_transport + import mcdc.transport as transport + + transport.simulation.source_loop = gpu_transport.simulation.source_loop + + +# ====================================================================================== +# Forward declaration +# ====================================================================================== + +# Main types +none_type = None +simulation_type = None +data_type = None + +# Access functions +state_spec = None +access_simulation = None +access_data_ptr = None +access_group = None +access_thread = None +particle_gpu = None +particle_record_gpu = None + +# Asynchronous transport kernels +step_async = None +find_cell_async = None + +# Memory allocations +alloc_managed_bytes = None +alloc_device_bytes = None + + +def forward_declare_gpu_program(): + import harmonize + import mcdc.numba_types as type_ + + # Get to set the globals + global none_type, simulation_type, data_type + global state_spec, access_simulation, access_data_ptr, access_group, access_thread, particle_gpu, particle_record_gpu + global step_async, find_cell_async + global alloc_managed_bytes, alloc_device_bytes # Compilation check if MPI.COMM_WORLD.Get_rank() == 0: @@ -29,10 +89,6 @@ def build_gpu_program(mcdc_container, data): else: harmonize.config.should_compile(harmonize.config.ShouldCompile.NEVER) - # ================================================================================== - # Forward declaration - # ================================================================================== - # ROCm and CUDA paths if config.args.gpu_cuda_path != None: harmonize.config.set_cuda_path(config.args.gpu_cuda_path) @@ -47,25 +103,25 @@ def build_gpu_program(mcdc_container, data): # Set access functions state_spec = ( { - "global": simulation_type, + "simulation": simulation_type, "data": data_type, }, none_type, none_type, ) access_fns = harmonize.RuntimeSpec.access_fns(state_spec) - simulation_gpu = access_fns["device"]["global"]["indirect"] - data_gpu = access_fns["device"]["data"]["direct"] - group_gpu = access_fns["group"] - thread_gpu = access_fns["thread"] + access_simulation = access_fns["device"]["simulation"]["indirect"] + access_data_ptr = access_fns["device"]["data"]["direct"] + access_group = access_fns["group"] + access_thread = access_fns["thread"] particle_gpu = nb.from_dtype(type_.particle) particle_record_gpu = nb.from_dtype(type_.particle_data) # Functions, and their signatures - def step(prog: nb.uintp, P: particle_gpu): + def step(program: nb.uintp, particle: particle_gpu): pass - def find_cell(prog: nb.uintp, P: particle_gpu): + def find_cell(program: nb.uintp, particle: particle_gpu): pass # Asynchronous versions @@ -75,52 +131,107 @@ def find_cell(prog: nb.uintp, P: particle_gpu): interface = harmonize.RuntimeSpec.program_interface() halt_early = interface["halt_early"] - # ================================================================================== - # TODO: "gpu_sources_spec" - # ================================================================================== + # Byte allocators + alloc_managed_bytes = harmonize.alloc_managed_bytes + alloc_device_bytes = harmonize.alloc_device_bytes + + +# ====================================================================================== +# Program builder +# ====================================================================================== + +alloc_state = None +free_state = None + +alloc_program = None +free_program = None + +load_state_device_simulation = None +store_state_device_simulation = None +store_pointer_state_device_simulation = None + +load_state_device_data = None +store_state_device_data = None +store_pointer_state_device_data = None + +init_program = None +exec_program = None +complete = None +clear_flags = None +set_device = None + +ARENA_SIZE = 0 +BLOCK_COUNT = 0 + + +def build_gpu_program(data_size): + import harmonize + import mcdc.numba_types as type_ + import mcdc.transport.util as util + + from mcdc.transport.simulation import generate_source_particle, step_particle + + global alloc_state, free_state + + global alloc_program, free_program + + global load_state_device_simulation, store_state_device_simulation, store_pointer_state_device_simulation + + global load_state_device_data, store_state_device_data, store_pointer_state_device_data + + global init_program, exec_program, complete, clear_flags, set_device + global ARENA_SIZE, BLOCK_COUNT + + shape = eval(f"{(data_size,)}") # ============== # Base functions # ============== - def make_work(prog: nb.uintp) -> nb.boolean: - mcdc = adapt.mcdc_global(prog) + def make_work(program: nb.uintp) -> nb.boolean: + simulation = access_simulation(program) + data_ptr = access_data_ptr(program) + data = harmonize.array_from_ptr(data_ptr, shape, nb.float64) - idx_work = adapt.global_add(mcdc["mpi_work_iter"], 0, 1) + util.atomic_add(simulation["mpi_work_iter"], 0, 1) + idx_work = simulation["mpi_work_iter"][0] - if idx_work >= mcdc["mpi_work_size"]: + if idx_work >= simulation["mpi_work_size"]: return False + work_start = simulation["mpi_work_start"] + generate_source_particle( - mcdc["mpi_work_start"], nb.uint64(idx_work), mcdc["source_seed"], prog + simulation["mpi_work_start"], + nb.uint64(idx_work), + simulation["source_seed"], + program, + data, ) return True - def initialize(prog: nb.uintp): + def initialize(program: nb.uintp): pass - def finalize(prog: nb.uintp): + def finalize(program: nb.uintp): pass # ================ # Async. functions # ================ - shape = eval(f"{adapt.tally_shape_literal}") - - def step(prog: nb.uintp, P_input: adapt.particle_gpu): - mcdc = adapt.mcdc_global(prog) - data_ptr = adapt.mcdc_data(prog) - data = adapt.harm.array_from_ptr(data_ptr, shape, nb.float64) - P_arr = adapt.local_array(1, type_.particle) - P_arr[0] = P_input - P = P_arr[0] - if P["fresh"]: - prep_particle(P_arr, prog) - P["fresh"] = False - step_particle(P_arr, data, prog) - if P["alive"]: - adapt.step_async(prog, P) + def step(program: nb.uintp, particle_input: particle_gpu): + simulation = access_simulation(program) + data_ptr = access_data_ptr(program) + data = harmonize.array_from_ptr(data_ptr, shape, nb.float64) + + particle_container = util.local_array(1, type_.particle) + particle_container[0] = particle_input + particle = particle_container[0] + particle["fresh"] = False + step_particle(particle_container, program, data) + if particle["alive"]: + step_async(program, particle) # Bind them all base_fns = (initialize, finalize, make_work) @@ -128,14 +239,7 @@ def step(prog: nb.uintp, P_input: adapt.particle_gpu): src_spec = harmonize.RuntimeSpec("mcdc_source", state_spec, base_fns, async_fns) harmonize.RuntimeSpec.bind_specs() - # ================================================================================== - # - # ================================================================================== - - rank = MPI.COMM_WORLD.Get_rank() - - MPI.COMM_WORLD.Barrier() - + # Load the specs harmonize.RuntimeSpec.load_specs() if config.args.gpu_strategy == "async": @@ -144,69 +248,99 @@ def step(prog: nb.uintp, P_input: adapt.particle_gpu): else: src_fns = src_spec.event_functions() - arena_size = config.args.gpu_arena_size - block_count = config.args.gpu_block_count + ARENA_SIZE = config.args.gpu_arena_size + BLOCK_COUNT = config.args.gpu_block_count alloc_state = src_fns["alloc_state"] free_state = src_fns["free_state"] - src_alloc_program = src_fns["alloc_program"] - src_free_program = src_fns["free_program"] - src_load_global = src_fns["load_state_device_global"] - src_store_global = src_fns["store_state_device_global"] - src_store_pointer_global = src_fns["store_pointer_state_device_global"] - src_load_data = src_fns["load_state_device_data"] - src_store_data = src_fns["store_state_device_data"] - src_store_pointer_data = src_fns["store_pointer_state_device_data"] - src_init_program = src_fns["init_program"] - src_exec_program = src_fns["exec_program"] - src_complete = src_fns["complete"] - src_clear_flags = src_fns["clear_flags"] - src_set_device = src_fns["set_device"] - - # ================================================================================== - # - # ================================================================================== - - """ - global loop_source - loop_source = gpu_loop_source - # - # Overwrite function - for impl in target_rosters["cpu"].values(): - overwrite_func(impl, impl) - """ - - # ================================================================================== - # Setup - # ================================================================================== - - device_id = rank % config.args.gpu_share_stride - - mcdc = mcdc_container[0] - src_set_device(device_id) - mcdc["gpu_meta"]["state_pointer"] = cast_voidptr_to_uintp(alloc_state()) + alloc_program = src_fns["alloc_program"] + free_program = src_fns["free_program"] + + load_state_device_simulation = src_fns["load_state_device_simulation"] + store_state_device_simulation = src_fns["store_state_device_simulation"] + store_pointer_state_device_simulation = src_fns[ + "store_pointer_state_device_simulation" + ] + + load_state_device_data = src_fns["load_state_device_data"] + store_state_device_data = src_fns["store_state_device_data"] + store_pointer_state_device_data = src_fns["store_pointer_state_device_data"] + + init_program = src_fns["init_program"] + exec_program = src_fns["exec_program"] + complete = src_fns["complete"] + clear_flags = src_fns["clear_flags"] + set_device = src_fns["set_device"] + + +# ====================================================================================== +# Setup GPU +# ====================================================================================== + +from numba import njit + +rank = MPI.COMM_WORLD.Get_rank() +device_id = rank % config.args.gpu_share_stride + + +@njit +def setup_gpu_program(simulation_container, data): + simulation = simulation_container[0] + + set_device(device_id) + simulation["gpu_meta"]["state_pointer"] = cast_voidptr_to_uintp(alloc_state()) + if config.gpu_state_storage == "separate": - src_store_pointer_global( - mcdc["gpu_meta"]["state_pointer"], mcdc["gpu_meta"]["global_pointer"] + store_pointer_state_device_simulation( + simulation["gpu_meta"]["state_pointer"], + simulation["gpu_meta"]["simulation_pointer"], ) - src_store_pointer_data( - mcdc["gpu_meta"]["state_pointer"], mcdc["gpu_meta"]["tally_pointer"] + store_pointer_state_device_data( + simulation["gpu_meta"]["state_pointer"], + simulation["gpu_meta"]["data_pointer"], ) else: - src_store_pointer_global(mcdc["gpu_meta"]["state_pointer"], mcdc_container) - src_store_pointer_data(mcdc["gpu_meta"]["state_pointer"], data) + store_pointer_state_device_simulation( + simulation["gpu_meta"]["state_pointer"], simulation_container + ) + store_pointer_state_device_data(simulation["gpu_meta"]["state_pointer"], data) - mcdc["gpu_meta"]["source_program_pointer"] = cast_voidptr_to_uintp( - src_alloc_program(mcdc["gpu_meta"]["state_pointer"], arena_size) + simulation["gpu_meta"]["program_pointer"] = cast_voidptr_to_uintp( + alloc_program(simulation["gpu_meta"]["state_pointer"], ARENA_SIZE) ) - src_init_program(mcdc["gpu_meta"]["source_program_pointer"], block_count) - return + init_program(simulation["gpu_meta"]["program_pointer"], BLOCK_COUNT) + + +@njit +def teardown_gpu_program(simulation): + free_program(cast_uintp_to_voidptr(simulation["gpu_meta"]["program_pointer"])) + free_state(cast_uintp_to_voidptr(simulation["gpu_meta"]["state_pointer"])) -def teardown_gpu_program(mcdc): - src_free_program(cast_uintp_to_voidptr(mcdc["gpu_meta"]["source_program_pointer"])) - free_state(cast_uintp_to_voidptr(mcdc["gpu_meta"]["state_pointer"])) +# ====================================================================================== +# Simulation structure and data creators +# ====================================================================================== + + +def create_data_array(size, dtype): + if config.gpu_state_storage == "managed": + data_tally_ptr = harmonize.alloc_managed_bytes(size) + else: + data_tally_ptr = harmonize.alloc_device_bytes(size) + data_tally_uint = cast_voidptr_to_uintp(data_tally_ptr) + data_tally = nb.carray(data_tally_ptr, (size,), dtype) + return data_tally, data_tally_uint + + +def create_mcdc_container(dtype): + if config.gpu_state_storage == "managed": + mcdc_ptr = harmonize.alloc_managed_bytes(dtype.itemsize) + else: + mcdc_ptr = harmonize.alloc_device_bytes(dtype.itemsize) + mcdc_uint = cast_voidptr_to_uintp(mcdc_ptr) + mcdc_container = nb.carray(mcdc_ptr, (1,), dtype) + return mcdc_container, mcdc_uint # ====================================================================================== diff --git a/mcdc/code_factory/gpu/transport/__init__.py b/mcdc/code_factory/gpu/transport/__init__.py new file mode 100644 index 00000000..6dfd8b41 --- /dev/null +++ b/mcdc/code_factory/gpu/transport/__init__.py @@ -0,0 +1,4 @@ +import mcdc.code_factory.gpu.transport.geometry as geometry +import mcdc.code_factory.gpu.transport.particle_bank as particle_bank +import mcdc.code_factory.gpu.transport.simulation as simulation +import mcdc.code_factory.gpu.transport.util as util diff --git a/mcdc/code_factory/gpu/transport/geometry/__init__.py b/mcdc/code_factory/gpu/transport/geometry/__init__.py new file mode 100644 index 00000000..e10e84a5 --- /dev/null +++ b/mcdc/code_factory/gpu/transport/geometry/__init__.py @@ -0,0 +1 @@ +import mcdc.code_factory.gpu.transport.geometry.interface as interface diff --git a/mcdc/code_factory/gpu/transport/geometry/interface.py b/mcdc/code_factory/gpu/transport/geometry/interface.py index 7b4effcc..e460f80e 100644 --- a/mcdc/code_factory/gpu/transport/geometry/interface.py +++ b/mcdc/code_factory/gpu/transport/geometry/interface.py @@ -6,6 +6,6 @@ @njit -def report_lost_particle(particle_container, mcdc): +def report_lost_particle(particle_container, simulation): particle = particle_container[0] particle["alive"] = False diff --git a/mcdc/code_factory/gpu/transport/particle_bank.py b/mcdc/code_factory/gpu/transport/particle_bank.py index c8133a89..aee6c20b 100644 --- a/mcdc/code_factory/gpu/transport/particle_bank.py +++ b/mcdc/code_factory/gpu/transport/particle_bank.py @@ -1,18 +1,31 @@ from numba import njit +### + +import mcdc.numba_types as type_ +import mcdc.transport.particle as particle_module +import mcdc.transport.util as util +import mcdc.code_factory.gpu.program_builder as gpu_program + +from mcdc.constant import GPU_ASYNC_SIMPLE + # ============================================================================= # Bank and pop particle # ============================================================================= @njit -def bank_active_particle(P_rec_arr, mcdc): - particle_container = local_array(1, type_.particle) - kernel.recordlike_to_particle(particle_container, P_rec_arr) - if SIMPLE_ASYNC: - step_async(prog, particle_container[0]) +def bank_active_particle(particle_container, program): + simulation = util.access_simulation(program) + + active_particle_container = util.local_array(1, type_.particle) + particle_module.copy(active_particle_container, particle_container) + if simulation["settings"]["gpu_async_type"] == GPU_ASYNC_SIMPLE: + gpu_program.step_async(program, active_particle_container[0]) + """ else: - find_cell_async(prog, particle_container[0]) + gpu_program.find_cell_async(program, active_particle_container[0]) + """ @njit diff --git a/mcdc/code_factory/gpu/transport/simulation.py b/mcdc/code_factory/gpu/transport/simulation.py index 0ce2a722..591ec618 100644 --- a/mcdc/code_factory/gpu/transport/simulation.py +++ b/mcdc/code_factory/gpu/transport/simulation.py @@ -1,109 +1,31 @@ -from mpi4py import MPI -import mcdc.code_factory.gpu.adapt as adapt +import harmonize -caching = config.caching - -# ============================================================================= -# Functions for GPU Interop -# ============================================================================= - -# The symbols declared below will be overwritten to reference external code that -# manages GPU execution (if GPU execution is supported and selected) -alloc_state, free_state = [None] * 2 - -src_alloc_program, src_free_program = [None] * 2 -( - src_load_global, - src_load_constant, - src_store_global, - src_store_data, - src_store_pointer_data, -) = [None] * 5 -src_init_program, src_exec_program, src_complete, src_clear_flags = [None] * 4 - -pre_alloc_program, pre_free_program = [None] * 2 -pre_load_global, pre_load_data, pre_store_global, pre_store_data = [None] * 4 -pre_init_program, pre_exec_program, pre_complete, pre_clear_flags = [None] * 4 - - -# If GPU execution is supported and selected, the functions shown below will -# be redefined to overwrite the above symbols and perform initialization/ -# finalization of GPU state -@njit -def setup_gpu(mcdc, data_tally): - pass - - -@njit -def teardown_gpu(mcdc): - pass - - -def gpu_sources_spec(): - def make_work(prog: nb.uintp) -> nb.boolean: - mcdc = adapt.mcdc_global(prog) - - atomic_add(mcdc["mpi_work_iter"], 0, 1) - idx_work = mcdc["mpi_work_iter"][0] - - if idx_work >= mcdc["mpi_work_size"]: - return False - - generate_source_particle( - mcdc["mpi_work_start"], nb.uint64(idx_work), mcdc["source_seed"], prog - ) - return True - - def initialize(prog: nb.uintp): - pass +from numba import njit - def finalize(prog: nb.uintp): - pass +### - base_fns = (initialize, finalize, make_work) +import mcdc.code_factory.gpu.program_builder as gpu_module +import mcdc.config as config +import mcdc.transport.particle_bank as particle_bank_module - def step(prog: nb.uintp, P_input: adapt.particle_gpu): - mcdc = adapt.mcdc_global(prog) - data = adapt.mcdc_data(prog) - particle_container = np.zeros(1, type_.particle) - particle_container[0] = P_input - particle = particle_container[0] - if particle["fresh"]: - prep_particle(particle_container, prog) - particle["fresh"] = False - step_particle(particle_container, data, prog) - if particle["alive"]: - adapt.step_async(prog, P) +from mcdc.constant import GPU_STORAGE_SEPARATE, GPU_STRATEGY_ASYNC +from mcdc.transport.simulation import source_closeout - async_fns = [step] - return adapt.harm.RuntimeSpec("mcdc_source", adapt.state_spec, base_fns, async_fns) - - -BLOCK_COUNT = config.args.gpu_block_count - -ASYNC_EXECUTION = config.args.gpu_strategy == "async" +caching = config.caching @njit(cache=caching) -def gpu_loop_source(seed, data, mcdc): - - # Progress bar indicator - N_prog = 0 - - if mcdc["technique"]["domain_decomposition"]: - particle_bank_module.dd_check_in(mcdc) - - # ===================================================================== - # GPU Interop - # ===================================================================== - +def source_loop(seed, simulation, data): # For async execution iter_count = 655360000 # For event-based execution - batch_size = 1 + batch_size = 64 - full_work_size = mcdc["mpi_work_size"] - if ASYNC_EXECUTION: + settings = simulation["settings"] + + full_work_size = simulation["mpi_work_size"] + + if settings["gpu_strategy"] == GPU_STRATEGY_ASYNC: phase_size = 1000000000 else: phase_size = 1000000 @@ -111,177 +33,53 @@ def gpu_loop_source(seed, data, mcdc): for phase in range(phase_count): - mcdc["mpi_work_iter"][0] = phase_size * phase - mcdc["mpi_work_size"] = min(phase_size * (phase + 1), full_work_size) - mcdc["source_seed"] = seed + simulation["mpi_work_iter"][0] = phase_size * phase + simulation["mpi_work_size"] = min(phase_size * (phase + 1), full_work_size) + simulation["source_seed"] = seed # Store the global state to the GPU - src_store_constant(mcdc["gpu_state_pointer"], mcdc) - src_store_data(mcdc["gpu_state_pointer"], data) + if settings["gpu_storage"] == GPU_STORAGE_SEPARATE: + harmonize.memcpy_host_to_device( + simulation["gpu_meta"]["state_pointer"], simulation + ) + harmonize.memcpy_host_to_device( + simulation["gpu_meta"]["state_pointer"], data + ) # Execute the program, and continue to do so until it is done - if ASYNC_EXECUTION: - src_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, iter_count) - while not src_complete(mcdc["source_program_pointer"]): - particle_bank_module.dd_particle_send(mcdc) - src_exec_program( - mcdc["source_program_pointer"], BLOCK_COUNT, iter_count + block_count = gpu_module.BLOCK_COUNT + + if settings["gpu_strategy"] == GPU_STRATEGY_ASYNC: + gpu_module.exec_program( + simulation["gpu_meta"]["program_pointer"], block_count, iter_count + ) + while not gpu_module.complete(simulation["gpu_meta"]["program_pointer"]): + gpu_module.exec_program( + simulation["gpu_meta"]["program_pointer"], block_count, iter_count ) else: - src_exec_program(mcdc["source_program_pointer"], BLOCK_COUNT, batch_size) - while not src_complete(mcdc["source_program_pointer"]): - particle_bank_module.dd_particle_send(mcdc) - src_exec_program( - mcdc["source_program_pointer"], BLOCK_COUNT, batch_size + gpu_module.exec_program( + simulation["gpu_meta"]["program_pointer"], block_count, batch_size + ) + while not gpu_module.complete(simulation["gpu_meta"]["program_pointer"]): + gpu_module.exec_program( + simulation["gpu_meta"]["program_pointer"], block_count, batch_size ) + gpu_module.clear_flags(simulation["gpu_meta"]["program_pointer"]) # Recover the original program state - src_load_constant(mcdc, mcdc["gpu_state_pointer"]) - src_load_data(data, mcdc["gpu_state_pointer"]) - src_clear_flags(mcdc["source_program_pointer"]) - - mcdc["mpi_work_size"] = full_work_size - - particle_bank_module.set_bank_size(mcdc["bank_active"], 0) - - # ===================================================================== - # Closeout (Moved out of the typical particle loop) - # ===================================================================== - - source_closeout(mcdc, 1, 1, data) - - if mcdc["technique"]["domain_decomposition"]: - source_dd_resolution(data, mcdc) - - -def build_gpu_progs(input_deck, args): - - STRAT = args.gpu_strategy - - src_spec = gpu_sources_spec() - - adapt.harm.RuntimeSpec.bind_specs() - - rank = MPI.COMM_WORLD.Get_rank() - device_id = rank % args.gpu_share_stride - - if MPI.COMM_WORLD.Get_size() > 1: - MPI.COMM_WORLD.Barrier() - - adapt.harm.RuntimeSpec.load_specs() - - if STRAT == "async": - args.gpu_arena_size = args.gpu_arena_size // 32 - src_fns = src_spec.async_functions() - pre_fns = pre_spec.async_functions() - else: - src_fns = src_spec.event_functions() - pre_fns = pre_spec.event_functions() - - ARENA_SIZE = args.gpu_arena_size - BLOCK_COUNT = args.gpu_block_count - - global alloc_state, free_state - alloc_state = src_fns["alloc_state"] - free_state = src_fns["free_state"] - - global src_alloc_program, src_free_program - global src_load_global, src_store_global, src_load_data, src_store_data, src_store_pointer_data - global src_init_program, src_exec_program, src_complete, src_clear_flags - src_alloc_program = src_fns["alloc_program"] - src_free_program = src_fns["free_program"] - src_load_global = src_fns["load_state_device_global"] - src_store_global = src_fns["store_state_device_global"] - src_store_pointer_global = src_fns["store_pointer_state_device_global"] - src_load_data = src_fns["load_state_device_data"] - src_store_data = src_fns["store_state_device_data"] - src_store_pointer_data = src_fns["store_pointer_state_device_data"] - src_init_program = src_fns["init_program"] - src_exec_program = src_fns["exec_program"] - src_complete = src_fns["complete"] - src_clear_flags = src_fns["clear_flags"] - src_set_device = src_fns["set_device"] - - global pre_alloc_program, pre_free_program - global pre_load_global, pre_store_global, pre_load_data, pre_store_data - global pre_init_program, pre_exec_program, pre_complete, pre_clear_flags - pre_alloc_state = pre_fns["alloc_state"] - pre_free_state = pre_fns["free_state"] - pre_alloc_program = pre_fns["alloc_program"] - pre_free_program = pre_fns["free_program"] - pre_load_global = pre_fns["load_state_device_global"] - pre_store_global = pre_fns["store_state_device_global"] - pre_load_data = pre_fns["load_state_device_data"] - pre_store_data = pre_fns["store_state_device_data"] - pre_init_program = pre_fns["init_program"] - pre_exec_program = pre_fns["exec_program"] - pre_complete = pre_fns["complete"] - pre_clear_flags = pre_fns["clear_flags"] - - @njit - def real_setup_gpu(mcdc_array, data_tally): - mcdc = mcdc_array[0] - src_set_device(device_id) - arena_size = ARENA_SIZE - mcdc["gpu_meta"]["state_pointer"] = adapt.cast_voidptr_to_uintp(alloc_state()) - # src_store_global(mcdc["gpu_meta"]["state_pointer"], mcdc_array[0]) if config.gpu_state_storage == "separate": - src_store_pointer_global( - mcdc["gpu_meta"]["state_pointer"], mcdc["gpu_meta"]["global_pointer"] + harmonize.memcpy_device_to_host( + simulation, simulation["gpu_meta"]["state_pointer"] ) - src_store_pointer_data( - mcdc["gpu_meta"]["state_pointer"], mcdc["gpu_meta"]["tally_pointer"] + harmonize.memcpy_device_to_host( + data, simulation["gpu_meta"]["state_pointer"] ) - else: - src_store_pointer_global(mcdc["gpu_meta"]["state_pointer"], mcdc_array) - src_store_pointer_data(mcdc["gpu_meta"]["state_pointer"], data_tally) - - mcdc["gpu_meta"]["source_program_pointer"] = adapt.cast_voidptr_to_uintp( - src_alloc_program(mcdc["gpu_meta"]["state_pointer"], ARENA_SIZE) - ) - src_init_program(mcdc["gpu_meta"]["source_program_pointer"], BLOCK_COUNT) - return - - @njit - def real_teardown_gpu(mcdc): - src_free_program( - adapt.cast_uintp_to_voidptr(mcdc["gpu_meta"]["source_program_pointer"]) - ) - free_state(adapt.cast_uintp_to_voidptr(mcdc["gpu_meta"]["state_pointer"])) - - global setup_gpu, teardown_gpu - setup_gpu = real_setup_gpu - teardown_gpu = real_teardown_gpu - - global loop_source - loop_source = gpu_loop_source - - -# ============================================================================= -# Functions for GPU Interop -# ============================================================================= - -# The symbols declared below will be overwritten to reference external code that -# manages GPU execution (if GPU execution is supported and selected) -alloc_state, free_state = [None] * 2 - -src_alloc_program, src_free_program = [None] * 2 -src_load_constant, src_load_constant, src_store_constant, src_store_data = [None] * 4 -src_init_program, src_exec_program, src_complete, src_clear_flags = [None] * 4 - -pre_alloc_program, pre_free_program = [None] * 2 -pre_load_constant, pre_load_data, pre_store_constant, pre_store_data = [None] * 4 -pre_init_program, pre_exec_program, pre_complete, pre_clear_flags = [None] * 4 + gpu_module.clear_flags(simulation["gpu_meta"]["program_pointer"]) -# If GPU execution is supported and selected, the functions shown below will -# be redefined to overwrite the above symbols and perform initialization/ -# finalization of GPU state -@njit -def setup_gpu(mcdc): - pass + simulation["mpi_work_size"] = full_work_size + particle_bank_module.set_bank_size(simulation["bank_active"], 0) -@njit -def teardown_gpu(mcdc): - pass + source_closeout(simulation, 1, 1, data) diff --git a/mcdc/code_factory/gpu/transport/util.py b/mcdc/code_factory/gpu/transport/util.py index 85ff979c..2acb013c 100644 --- a/mcdc/code_factory/gpu/transport/util.py +++ b/mcdc/code_factory/gpu/transport/util.py @@ -1,6 +1,165 @@ -from numba import njit +import harmonize +import numba as nb +import numpy as np + +from numba import njit, types @njit def atomic_add(array, idx, value): - return harmonize.array_atomic_add(array, idx, value) + harmonize.array_atomic_add(array, idx, value) + + +# ============================================================================= +# Generic GPU/CPU local array variable constructors +# ============================================================================= + + +def local_array(shape, dtype): + return np.zeros(shape, dtype=dtype) + + +@nb.extending.type_callable(local_array) +def type_local_array(context): + + from numba.core.typing.npydecl import parse_dtype, parse_shape + + if isinstance(context, nb.core.typing.context.Context): + + # Function repurposed from Numba's ol_np_empty. + def typer(shape, dtype): + nb.np.arrayobj._check_const_str_dtype("empty", dtype) + + # Only integer literals and tuples of integer literals are valid + # shapes + if isinstance(shape, types.Integer): + if not isinstance(shape, types.IntegerLiteral): + raise nb.core.errors.UnsupportedError( + f"Integer shape type {shape} is not literal." + ) + elif isinstance(shape, (types.Tuple, types.UniTuple)): + if any([not isinstance(s, types.IntegerLiteral) for s in shape]): + raise nb.core.errors.UnsupportedError( + f"At least one element of shape tuple type{shape} is not an integer literal." + ) + else: + raise nb.core.errors.UnsupportedError( + f"Shape is of unsupported type {shape}." + ) + + # No default arguments. + nb_dtype = parse_dtype(dtype) + nb_shape = parse_shape(shape) + + if nb_dtype is not None and nb_shape is not None: + retty = types.Array(dtype=nb_dtype, ndim=nb_shape, layout="C") + # Inlining the signature construction from numpy_empty_nd + sig = retty(shape, dtype) + return sig + else: + msg = f"Cannot parse input types to function np.empty({shape}, {dtype})" + raise nb.errors.TypingError(msg) + + return typer + + elif isinstance(context, nb.cuda.target.CUDATypingContext): + + # Function repurposed from Numba's Cuda_array_decl. + def typer(shape, dtype): + + # Only integer literals and tuples of integer literals are valid + # shapes + if isinstance(shape, types.Integer): + if not isinstance(shape, types.IntegerLiteral): + return None + elif isinstance(shape, (types.Tuple, types.UniTuple)): + if any([not isinstance(s, types.IntegerLiteral) for s in shape]): + return None + else: + return None + + ndim = parse_shape(shape) + nb_dtype = parse_dtype(dtype) + if nb_dtype is not None and ndim is not None: + return types.Array(dtype=nb_dtype, ndim=ndim, layout="C") + + return typer + + elif isinstance(context, nb.hip.target.HIPTypingContext): + + def typer(shape, dtype): + # Only integer literals and tuples of integer literals are valid + # shapes + if isinstance(shape, types.Integer): + if not isinstance(shape, types.IntegerLiteral): + return None + elif isinstance(shape, (types.Tuple, types.UniTuple)): + if any([not isinstance(s, types.IntegerLiteral) for s in shape]): + return None + else: + return None + + ndim = parse_shape(shape) + nb_dtype = parse_dtype(dtype) + if nb_dtype is not None and ndim is not None: + result = types.Array(dtype=nb_dtype, ndim=ndim, layout="C") + return result + + return typer + + else: + raise nb.core.errors.UnsupportedError(f"Unsupported target context {context}.") + + +@nb.extending.lower_builtin(local_array, types.IntegerLiteral, types.Any) +def builtin_local_array(context, builder, sig, args): + + shape, dtype = sig.args + + from numba.core.typing.npydecl import parse_dtype, parse_shape + import numba.np.arrayobj as arrayobj + + if isinstance(context, nb.core.cpu.CPUContext): + + # No default arguments. + nb_dtype = parse_dtype(dtype) + nb_shape = parse_shape(shape) + + retty = types.Array(dtype=nb_dtype, ndim=nb_shape, layout="C") + + # In ol_np_empty, the reference type of the array is fed into the + # signatrue as a third argument. This third argument is not used by + # _parse_empty_args. + sig = retty(shape, dtype) + + arrtype, shapes = arrayobj._parse_empty_args(context, builder, sig, args) + ary = arrayobj._empty_nd_impl(context, builder, arrtype, shapes) + + return ary._getvalue() + elif isinstance(context, nb.cuda.target.CUDATargetContext): + length = sig.args[0].literal_value + dtype = parse_dtype(sig.args[1]) + return nb.cuda.cudaimpl._generic_array( + context, + builder, + shape=(length,), + dtype=dtype, + symbol_name="_cudapy_harm_lmem", + addrspace=nb.cuda.cudadrv.nvvm.ADDRSPACE_LOCAL, + can_dynsized=False, + ) + elif isinstance(context, nb.hip.target.HIPTargetContext): + length = sig.args[0].literal_value + dtype = parse_dtype(sig.args[1]) + result = nb.hip.typing_lowering.hip.lowering._generic_array( + context, + builder, + shape=(length,), + dtype=dtype, + symbol_name="_HIPpy_lmem", + addrspace=nb.hip.amdgcn.ADDRSPACE_LOCAL, + can_dynsized=False, + ) + return result + else: + raise nb.core.errors.UnsupportedError(f"Unsupported target context {context}.") diff --git a/mcdc/code_factory/literals_generator.py b/mcdc/code_factory/literals_generator.py new file mode 100644 index 00000000..3e5329a5 --- /dev/null +++ b/mcdc/code_factory/literals_generator.py @@ -0,0 +1,25 @@ +import numpy as np + +from numba import njit + + +def _literalize(value): + namespace = {} + jit_str = f"@njit\ndef impl():\n return {value}\n" + exec(jit_str, globals(), namespace) + return namespace["impl"] + + +def make_literals(simulation): + import mcdc.literals as literals + + # RPN evaluation buffer size + if len(simulation.cells) == 0: + rpn_evaluation_buffer_size = 1 + else: + rpn_evaluation_buffer_size = int( + max( + [np.sum(np.array(x.region_RPN_tokens) >= 0.0) for x in simulation.cells] + ) + ) + literals.rpn_evaluation_buffer_size = _literalize(rpn_evaluation_buffer_size) diff --git a/mcdc/code_factory/numba_objects_generator.py b/mcdc/code_factory/numba_objects_generator.py index d8648c0c..25653afb 100644 --- a/mcdc/code_factory/numba_objects_generator.py +++ b/mcdc/code_factory/numba_objects_generator.py @@ -2,14 +2,19 @@ #### +import importlib +import numba as nb import numpy as np -from pathlib import Path from mpi4py import MPI +from numba import njit +from numba.extending import intrinsic +from pathlib import Path #### import mcdc +import mcdc.code_factory.gpu.program_builder as gpu_builder import mcdc.config as config import mcdc.object_ as object_module import mcdc.object_.base as base @@ -26,15 +31,16 @@ from mcdc.util import flatten type_map = { - bool: "?", - float: "f8", - int: "i8", + bool: np.bool_, + float: np.float64, + int: np.int64, str: "U32", - np.bool_: "?", - np.float64: "f8", - np.int64: "i8", - np.uint64: "u8", + np.bool_: np.bool_, + np.float64: np.float64, + np.int64: np.int64, + np.uint64: np.uint64, np.str_: "U32", + np.uintp: np.uintp, } bank_names = ["bank_active", "bank_census", "bank_source", "bank_future"] @@ -200,21 +206,21 @@ def generate_numba_objects(simulation): # Add ID for non-singleton for class_ in mcdc_classes: if issubclass(class_, ObjectNonSingleton): - structures[class_.label].append(("ID", "i8")) + structures[class_.label].append(("ID", type_map[int])) # Set parent and child ID and type if polymorphic if issubclass(class_, ObjectPolymorphic): if class_.__name__[-4:] == "Base" or class_.__name__ == "Tally": - structures[class_.label].append(("child_type", "i8")) - structures[class_.label].append(("child_ID", "i8")) + structures[class_.label].append(("child_type", type_map[int])) + structures[class_.label].append(("child_ID", type_map[int])) else: - structures[class_.label].append(("parent_ID", "i8")) + structures[class_.label].append(("parent_ID", type_map[int])) - # Add particles to particle banks and add particle banks to the simulation + # Add particle data to particle banks and add particle banks to the simulation for name in bank_names: bank = getattr(simulation, name) size = int(bank.size[0]) structures[name] += [ - ("particles", into_dtype(structures["particle_data"]), (size,)) + ("particle_data", into_dtype(structures["particle_data"]), (size,)) ] # structures["simulation"] = [(name, into_dtype(structures[name]))] + structures[ @@ -222,7 +228,7 @@ def generate_numba_objects(simulation): ] # ================================================================================== - # Set records and data based on the simulation structures and objects + # Set records and data size based on the simulation structures and objects # ================================================================================== # Allocate object containers @@ -250,27 +256,11 @@ def generate_numba_objects(simulation): if type(item) in mcdc_classes: objects.append(item) - # Set the objects + # Set the records and the data size for object_ in objects: set_object(object_, annotations, structures, records, data) set_object(simulation, annotations, structures, records, data) - # Allocate the flattened data and re-set the objects - data["array"], data["pointer"] = create_data_array(data["size"], type_map[float]) - - data["size"] = 0 - records = {} - for mcdc_class in mcdc_classes: - if issubclass(mcdc_class, ObjectNonSingleton): - records[mcdc_class.label] = [] - else: - records[mcdc_class.label] = {} - records["simulation"] = records.pop("simulation") - - for object_ in objects: - set_object(object_, annotations, structures, records, data, set_data=True) - set_object(simulation, annotations, structures, records, data, set_data=True) - # ================================================================================== # Finalize the simulation object structure and set record # ================================================================================== @@ -291,7 +281,7 @@ def generate_numba_objects(simulation): new_structure.append( (field, into_dtype(structures[item[2].label]), (N,)) ) - new_structure.append((f"N_{plural_to_singular(field)}", "i8")) + new_structure.append((f"N_{plural_to_singular(field)}", type_map[int])) record[f"N_{plural_to_singular(field)}"] = N # List of polymorphics @@ -306,7 +296,7 @@ def generate_numba_objects(simulation): (N,), ) ) - new_structure.append((f"N_{class_.label}", "i8")) + new_structure.append((f"N_{class_.label}", type_map[int])) record[f"N_{class_.label}"] = N # Singleton @@ -322,6 +312,12 @@ def generate_numba_objects(simulation): if MPI.COMM_WORLD.Get_rank() == 0: with open(f"{Path(mcdc.__file__).parent}/numba_types.py", "w") as f: text = "# The following is automatically generated by code_factory.py\n\n" + text += "from numpy import bool_\n" + text += "from numpy import float64\n" + text += "from numpy import int64\n" + text += "from numpy import uint64\n" + text += "from numpy import uintp\n" + text += "\n###\n\n" text += ( "from mcdc.code_factory.numba_objects_generator import into_dtype\n\n" ) @@ -329,9 +325,31 @@ def generate_numba_objects(simulation): for label in structures.keys(): text += f"{label} = into_dtype([\n" structure = structures[label] + + # GPU meta override + if label == "gpu_meta": + for item in structure: + if type(item[1]) != np.dtypes.VoidDType: + if isinstance(item[1], str): + dtype = f"'{item[1]}'" + elif item[1].__name__ == "uint64": + dtype = "uintp" + else: + dtype = item[1].__name__ + text += f" ('{item[0]}', {dtype}),\n" + text += "])\n\n" + continue + for item in structure: if type(item[1]) != np.dtypes.VoidDType: - text += f" {item},\n" + if isinstance(item[1], str): + dtype = f"'{item[1]}'" + else: + dtype = item[1].__name__ + if len(item) == 3: + text += f" ('{item[0]}', {dtype}, {item[2]}),\n" + else: + text += f" ('{item[0]}', {dtype}),\n" else: if len(item) == 3: text += f" ('{item[0]}', {plural_to_singular(item[0])}, {item[2]}),\n" @@ -341,15 +359,35 @@ def generate_numba_objects(simulation): f.write(text) + # ================================================================================== + # GPU preparation: Adapt transport functions, forward declare, and build program + # ================================================================================== + + if config.target == "gpu": + gpu_builder.forward_declare_gpu_program() + gpu_builder.adapt_transport_functions() + gpu_builder.build_gpu_program(data["size"]) + + # ================================================================================== + # Allocate the flattened data and re-set the objects + # ================================================================================== + + data["array"], data["pointer"] = create_data_array(data["size"]) + + data["size"] = 0 + for object_ in objects: + set_object(object_, annotations, structures, records, data, set_data=True) + set_object(simulation, annotations, structures, records, data, set_data=True) + # ================================================================================== # Set with records # ================================================================================== # The global structure/variable container - mcdc_simulation_arr, mcdc_simulation_pointer = create_mcdc_array( + mcdc_simulation_container, mcdc_simulation_pointer = create_simulation_container( into_dtype(structures["simulation"]) ) - mcdc_simulation = mcdc_simulation_arr[0] + mcdc_simulation = mcdc_simulation_container[0] record = records["simulation"] structure = structures["simulation"] @@ -387,7 +425,12 @@ def generate_numba_objects(simulation): for name in bank_names: mcdc_simulation[name]["tag"] = getattr(simulation, name).tag - return mcdc_simulation_arr, data["array"] + # GPU program setup + if config.target == "gpu": + gpu_builder.setup_gpu_program(mcdc_simulation_container, data["array"]) + gpu_builder.adapt_transport_functions_post_setup() + + return mcdc_simulation_container, data["array"] def set_structure(label, structures, accessor_targets, annotations): @@ -463,8 +506,8 @@ def polymorphic_base(x): elif simple_scalar: structure.append((field, type_map[hint])) elif simple_list or numpy_array: - structure.append((f"{field}_offset", "i8")) - structure.append((f"{field}_length", "i8")) + structure.append((f"{field}_offset", type_map[int])) + structure.append((f"{field}_length", type_map[int])) if hint_origin_shape is not None: accessor_target.append((f"{field}", hint_origin_shape)) else: @@ -472,13 +515,13 @@ def polymorphic_base(x): # MC/DC classes elif non_polymorphic(hint) or polymorphic_base(hint): - structure.append((f"{field}_ID", "i8")) + structure.append((f"{field}_ID", type_map[int])) # List of MC/DC classes elif list_of_non_polymorphics or list_of_polymorphic_bases: singular = plural_to_singular(field) - structure.append((f"N_{singular}", "i8")) - structure.append((f"{singular}_IDs_offset", "i8")) + structure.append((f"N_{singular}", type_map[int])) + structure.append((f"{singular}_IDs_offset", type_map[int])) if hint_origin_shape is not None: accessor_target.append((f"{singular}_IDs", hint_origin_shape)) else: @@ -640,42 +683,75 @@ def set_object( # ============================================================================= -# Global GPU/CPU Array Variable Constructors +# Global GPU/CPU variable array constructors # ============================================================================= -def create_data_array(size, dtype): - if config.target == "gpu": - import mcdc.code_factory.gpu.adapt as adapt - import harmonize, numba +def create_data_array(size): + if not config.target == "gpu": + data = np.zeros(size, dtype=np.float64) + return data, 0 + else: + return create_data_array_on_gpu(size * 8) - if config.gpu_state_storage == "managed": - data_tally_ptr = harmonize.alloc_managed_bytes(size) - else: - data_tally_ptr = harmonize.alloc_device_bytes(size) - data_tally_uint = adapt.voidptr_to_uintp(data_tally_ptr) - data_tally = numba.carray(data_tally_ptr, (size,), dtype) - return data_tally, data_tally_uint + +@njit +def create_data_array_on_gpu(size): + if config.gpu_state_storage == "managed": + data_ptr = gpu_builder.alloc_managed_bytes(size) else: - data_tally = np.zeros(size, dtype=dtype) - return data_tally, 0 + data_ptr = gpu_builder.alloc_device_bytes(size) + data_uint = voidptr_to_uintp(data_ptr) + data = nb.carray(data_ptr, (size,), dtype=np.float64) + return data, data_uint -def create_mcdc_array(dtype): - if config.target == "gpu": - import mcdc.code_factory.gpu.adapt as adapt - import harmonize, numba +def create_simulation_container(dtype): + if not config.target == "gpu": + simulation_container = np.zeros((1,), dtype=dtype) + return simulation_container, 0 + else: + return create_simulation_container_on_gpu(dtype, dtype.itemsize) - if config.gpu_state_storage == "managed": - mcdc_ptr = harmonize.alloc_managed_bytes(dtype.itemsize) - else: - mcdc_ptr = harmonize.alloc_device_bytes(dtype.itemsize) - mcdc_uint = adapt.voidptr_to_uintp(mcdc_ptr) - mcdc_array = numba.carray(mcdc_ptr, (1,), dtype) - return mcdc_array, mcdc_uint + +@njit +def create_simulation_container_on_gpu(dtype, size): + if config.gpu_state_storage == "managed": + simulation_ptr = gpu_builder.alloc_managed_bytes(size) else: - mcdc_array = np.zeros((1,), dtype=dtype) - return mcdc_array, 0 + simulation_ptr = gpu_builder.alloc_device_bytes(size) + simulation_uint = voidptr_to_uintp(simulation_ptr) + simulation = nb.carray(simulation_ptr, (1,), dtype) + return simulation, simulation_uint + + +# ============================================================================= +# Type casters +# ============================================================================= + + +@intrinsic +def cast_voidptr_to_uintp(typingctx, src): + # check for accepted types + if isinstance(src, nb.types.RawPointer): + # create the expected type signature + result_type = nb.types.uintp + sig = result_type(nb.types.voidptr) + + # defines the custom code generation + def codegen(context, builder, signature, args): + # llvm IRBuilder code here + [src] = args + rtype = signature.return_type + llrtype = context.get_value_type(rtype) + return builder.ptrtoint(src, llrtype) + + return sig, codegen + + +@njit() +def voidptr_to_uintp(value): + return cast_voidptr_to_uintp(value) # ====================================================================================== @@ -1130,49 +1206,3 @@ def singular_to_plural(word: str) -> str: parts[-1] = w + "s" return "_".join(parts) - - -# ============================================================================== -# MC/DC Member Array Sizes -# ============================================================================== - - -def literalize(value): - jit_str = f"@njit\ndef impl():\n return {value}\n" - exec(jit_str, globals(), locals()) - return eval("impl") - - -def rpn_buffer_size(): - pass - - -def make_size_rpn(cells): - global rpn_buffer_size - size = max([np.sum(np.array(x.region_RPN_tokens) >= 0.0) for x in cells]) - rpn_buffer_size = literalize(size) - - -# ====================================================================================== -# Make literals -# ====================================================================================== - - -def make_literals(simulation): - # Sizes - if len(simulation.cells) == 0: - rpn_evaluation_buffer_size = 1 - else: - rpn_evaluation_buffer_size = int( - max( - [np.sum(np.array(x.region_RPN_tokens) >= 0.0) for x in simulation.cells] - ) - ) - - path = f"{Path(mcdc.__file__).parent}" - with open(f"{path}/transport/literals.py", "w") as f: - text = "# The following is automatically generated by code_factory.py\n\n" - - text += f"rpn_evaluation_buffer_size = {rpn_evaluation_buffer_size}\n" - - f.write(text) diff --git a/mcdc/constant.py b/mcdc/constant.py index 29a487c8..0b7130f1 100644 --- a/mcdc/constant.py +++ b/mcdc/constant.py @@ -195,3 +195,20 @@ # Weight Windows Modifications WW_MIN = 0 WW_WOLLABER = 1 + +# ====================================================================================== +# GPU settings +# ====================================================================================== + +# GPU strategies +GPU_STRATEGY_SIMPLE_ASYNC = 0 +GPU_STRATEGY_ASYNC = 0 +GPU_STRATEGY_EVENT = 1 + +# GPU async. types +GPU_ASYNC_SIMPLE = 0 + +# GPU storage types +GPU_STORAGE_SEPARATE = 0 +GPU_STORAGE_MANAGED = 1 +GPU_STORAGE_UNITED = 2 diff --git a/mcdc/literals.py b/mcdc/literals.py new file mode 100644 index 00000000..3cec7f37 --- /dev/null +++ b/mcdc/literals.py @@ -0,0 +1,6 @@ +# The following will be replaced with their respective literals by +# code_factory/literals_generator.py + + +def rpn_evaluation_buffer_size(): + pass diff --git a/mcdc/main.py b/mcdc/main.py index 3f7631ce..1e9ff34a 100644 --- a/mcdc/main.py +++ b/mcdc/main.py @@ -21,23 +21,11 @@ def run(): time_total_start = MPI.Wtime() # Get settings and MPI master status - from mcdc.object_.simulation import simulation + from mcdc.object_.simulation import simulation as simulationPy - settings = simulation.settings + settings = simulationPy.settings master = MPI.COMM_WORLD.Get_rank() == 0 - # Override settings with command-line arguments - import mcdc.config as config - - if config.args.N_particle is not None: - settings.N_particle = config.args.N_particle - if config.args.N_batch is not None: - settings.N_batch = config.args.N_batch - if config.args.output is not None: - settings.output_name = config.args.output - if config.args.progress_bar is not None: - settings.use_progress_bar = config.args.progress_bar - # ================================================================================== # Preparation # ================================================================================== @@ -45,23 +33,27 @@ def run(): # TIMER: preparation time_prep_start = MPI.Wtime() + # Override settings with command-line arguments + override_settings() + # Generate the program state: - # - `mcdc`: the simulation structure, storing fixed side data and meta data + # - `simulation`: the simulation, storing fixed side data and meta data that + # describes arbitrarily-sized data # - `data`: a long 1D array storing arbitrarily-sized data of the simulation - # NOTE: The simulation structure needs to be generated in a container, which is a + # NOTE: The simulation structure to be generated in a container, which is a # a one-sized array that stores the structure. The container is needed to # ensure proper mutability and tracking of the structure when running in - # different kinds of machines supported. - mcdc_container, data = preparation() - mcdc = mcdc_container[0] + # different kinds of machines supported by the Numba compilation framework. + simulation_container, data = preparation() + simulation = simulation_container[0] # Print headers if master: print_module.print_banner() print_module.print_configuration() print(" Now running the particle transport...") - if settings.eigenvalue_mode: - print_module.print_eigenvalue_header(mcdc) + if settings.neutron_eigenvalue_mode: + print_module.print_eigenvalue_header(simulation) # TIMER: preparation time_prep_end = MPI.Wtime() @@ -76,10 +68,10 @@ def run(): # Run simulation import mcdc.transport.simulation as simulation_module - if settings.eigenvalue_mode: - simulation_module.eigenvalue_simulation(mcdc_container, data) + if settings.neutron_eigenvalue_mode: + simulation_module.eigenvalue_simulation(simulation_container, data) else: - simulation_module.fixed_source_simulation(mcdc_container, data) + simulation_module.fixed_source_simulation(simulation_container, data) # TIMER: simulation time_simulation_end = MPI.Wtime() @@ -94,7 +86,7 @@ def run(): time_output_start = MPI.Wtime() # Generate hdf5 output file - output_module.generate_output(mcdc, data) + output_module.generate_output(simulation, data) # TIMER: output time_output_end = MPI.Wtime() @@ -106,23 +98,19 @@ def run(): time_total_end = MPI.Wtime() # Manage timers - mcdc["runtime_total"] = time_total_end - time_total_start - mcdc["runtime_preparation"] = time_prep_end - time_prep_start - mcdc["runtime_simulation"] = time_simulation_end - time_simulation_start - mcdc["runtime_output"] = time_output_end - time_output_start - output_module.create_runtime_datasets(mcdc) + simulation["runtime_total"] = time_total_end - time_total_start + simulation["runtime_preparation"] = time_prep_end - time_prep_start + simulation["runtime_simulation"] = time_simulation_end - time_simulation_start + simulation["runtime_output"] = time_output_end - time_output_start + output_module.create_runtime_datasets(simulation) if master: - print_module.print_runtime(mcdc) + print_module.print_runtime(simulation) # ================================================================================== # Finalizing # ================================================================================== - # GPU teardowns if needed - if config.target == "gpu": - from mcdc.code_factory.gpu.program_builder import teardown_gpu_program - - teardown_gpu_program(mcdc) + finalize(simulation) # ====================================================================================== @@ -135,7 +123,7 @@ def preparation(): from mpi4py import MPI - from mcdc.object_.simulation import simulation + from mcdc.object_.simulation import simulation as simulationPy from mcdc.object_.material import ( Material, MaterialMG, @@ -149,10 +137,19 @@ def preparation(): # ================================================================================== # Get settings - settings = simulation.settings + settings = simulationPy.settings + + # Set appropriate time boundary + settings.time_boundary = min( + [settings.time_boundary] + [tally.time[-1] for tally in simulationPy.tallies] + ) + + # ================================================================================== + # Set material data as needed + # ================================================================================== # Set material compositions based on transported particles - for material in simulation.materials: + for material in simulationPy.materials: if not isinstance(material, Material): continue @@ -164,28 +161,25 @@ def preparation(): # Set nuclear and atomic data for transported particles if settings.neutron_transport: - for nuclide in simulation.nuclides: + for nuclide in simulationPy.nuclides: nuclide.set_neutron_data() - for material in simulation.materials: + for material in simulationPy.materials: if isinstance(material, Material): update_fissionable_from_nuclides(material) if settings.electron_transport: - for element in simulation.elements: + for element in simulationPy.elements: element.set_electron_data() # Set physics mode - if len(simulation.materials) == 0: + if len(simulationPy.materials) == 0: # Default physics in dummy mode - settings.multigroup_mode = True + settings.neutron_multigroup_mode = True else: - settings.multigroup_mode = isinstance(simulation.materials[0], MaterialMG) - - # Set appropriate time boundary - settings.time_boundary = min( - [settings.time_boundary] + [tally.time[-1] for tally in simulation.tallies] - ) + settings.neutron_multigroup_mode = isinstance( + simulationPy.materials[0], MaterialMG + ) # ================================================================================== # Adjust simulation parameters as needed @@ -194,29 +188,29 @@ def preparation(): # Reset time grid size of all tallies if census-based tally is desired if settings.use_census_based_tally: N_bin = settings.census_tally_frequency - for tally in simulation.tallies: + for tally in simulationPy.tallies: tally._use_census_based_tally(N_bin) # Normalize source probability norm = 0.0 - for source in simulation.sources: + for source in simulationPy.sources: norm += source.probability - for source in simulation.sources: + for source in simulationPy.sources: source.probability /= norm # Create root universe if not defined - if len(simulation.universes[0].cells) == 0: - simulation.universes[0].cells = simulation.cells + if len(simulationPy.universes[0].cells) == 0: + simulationPy.universes[0].cells = simulationPy.cells # Initial guess - simulation.k_eff = settings.k_init + simulationPy.k_eff = settings.k_init # Activate tally scoring for fixed-source - if not settings.eigenvalue_mode: - simulation.cycle_active = True + if not settings.neutron_eigenvalue_mode: + simulationPy.cycle_active = True # All active eigenvalue cycle? elif settings.N_inactive == 0: - simulation.cycle_active = True + simulationPy.cycle_active = True # ================================================================================== # Set particle bank sizes @@ -228,9 +222,9 @@ def preparation(): N_census = settings.N_census # Determine bank size - if settings.eigenvalue_mode or N_census == 1: + if settings.neutron_eigenvalue_mode or N_census == 1: settings.future_bank_buffer_ratio = 0.0 - if not settings.eigenvalue_mode and N_census == 1: + if not settings.neutron_eigenvalue_mode and N_census == 1: settings.census_bank_buffer_ratio = 0.0 settings.source_bank_buffer_ratio = 0.0 size_active = settings.active_bank_buffer @@ -239,23 +233,22 @@ def preparation(): size_future = int((settings.future_bank_buffer_ratio) * N_work) # Set bank size - simulation.bank_active.size[0] = size_active - simulation.bank_census.size[0] = size_census - simulation.bank_source.size[0] = size_source - simulation.bank_future.size[0] = size_future + simulationPy.bank_active.size[0] = size_active + simulationPy.bank_census.size[0] = size_census + simulationPy.bank_source.size[0] = size_source + simulationPy.bank_future.size[0] = size_future # ================================================================================== # Generate Numba-supported "Objects" # ================================================================================== from mcdc.code_factory.numba_objects_generator import generate_numba_objects + from mcdc.code_factory.literals_generator import make_literals - if MPI.COMM_WORLD.Get_rank() == 0: - from mcdc.code_factory.numba_objects_generator import make_literals + make_literals(simulationPy) - make_literals(simulation) - mcdc_container, data = generate_numba_objects(simulation) - mcdc = mcdc_container[0] + simulation_container, data = generate_numba_objects(simulationPy) + simulation = simulation_container[0] # Reload mcdc getters and setters import importlib @@ -272,7 +265,7 @@ def preparation(): # Pick physics model import mcdc.transport.physics as physics - if settings.multigroup_mode: + if settings.neutron_multigroup_mode: physics.neutron.particle_speed = physics.neutron.multigroup.particle_speed physics.neutron.macro_xs = physics.neutron.multigroup.macro_xs physics.neutron.neutron_production_xs = ( @@ -299,38 +292,81 @@ def preparation(): import h5py # All ranks, take turn - for i in range(mcdc["mpi_size"]): - if mcdc["mpi_rank"] == i: + for i in range(simulation["mpi_size"]): + if simulation["mpi_rank"] == i: if settings.use_source_file: with h5py.File(settings.source_file_name, "r") as f: # Get source particle size N_particle = f["particles_size"][()] # Redistribute work - mpi.distribute_work(N_particle, mcdc) - N_local = mcdc["mpi_work_size"] - start = mcdc["mpi_work_start"] + mpi.distribute_work(N_particle, simulation) + N_local = simulation["mpi_work_size"] + start = simulation["mpi_work_start"] end = start + N_local # Add particles to source bank - mcdc["bank_source"]["particles"][:N_local] = f["particles"][ + simulation["bank_source"]["particles"][:N_local] = f["particles"][ start:end ] - mcdc["bank_source"]["size"] = N_local + simulation["bank_source"]["size"] = N_local MPI.COMM_WORLD.Barrier() # ================================================================================== - # Platform targeting, adapters, and toggles for portability + # Finalize # ================================================================================== - # Build GPU program if desired + return simulation_container, data + + +# ====================================================================================== +# Misc. +# ====================================================================================== + + +def override_settings(): + import mcdc.config as config + from mcdc.object_.simulation import simulation as simulationPy + + settings = simulationPy.settings + + if config.args.N_particle is not None: + settings.N_particle = config.args.N_particle + if config.args.N_batch is not None: + settings.N_batch = config.args.N_batch + if config.args.output is not None: + settings.output_name = config.args.output + if config.args.progress_bar is not None: + settings.use_progress_bar = config.args.progress_bar + + # GPU settings if config.target == "gpu": - from mcdc.code_factory.gpu.program_builder import build_gpu_program + from mcdc.constant import ( + GPU_STRATEGY_ASYNC, + GPU_STRATEGY_EVENT, + GPU_STORAGE_SEPARATE, + GPU_STORAGE_MANAGED, + GPU_STORAGE_UNITED, + ) - build_gpu_program(mcdc_container, data) + if config.args.gpu_strategy == "async": + settings.gpu_strategy = GPU_STRATEGY_ASYNC + elif config.args.gpu_strategy == "event": + settings.gpu_strategy = GPU_STRATEGY_EVENT - # ================================================================================== - # Finalize - # ================================================================================== + if config.args.gpu_state_storage == "separate": + settings.gpu_storage = GPU_STORAGE_SEPARATE + elif config.args.gpu_state_storage == "managed": + settings.gpu_storage = GPU_STORAGE_MANAGED + elif config.args.gpu_state_storage == "united": + settings.gpu_storage = GPU_STORAGE_UNITED + + +def finalize(simulation): + import mcdc.config as config + + # GPU teardowns if needed + if config.target == "gpu": + from mcdc.code_factory.gpu.program_builder import teardown_gpu_program - return mcdc_container, data + teardown_gpu_program(simulation) diff --git a/mcdc/object_/gpu_tools.py b/mcdc/object_/gpu_tools.py index a7ad7d3c..b4d4965a 100644 --- a/mcdc/object_/gpu_tools.py +++ b/mcdc/object_/gpu_tools.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from numpy import uint64 +from numpy import uintp #### @@ -11,8 +11,9 @@ class GPUMeta(ObjectSingleton): # Annotations for Numba mode label: str = "gpu_meta" # - state_pointer: uint64 = uint64(0) - source_program_pointer: uint64 = uint64(0) - precursor_program_pointer: uint64 = uint64(0) - structure_pointer: uint64 = uint64(0) - data_pointer: uint64 = uint64(0) + state_pointer: uintp = uintp(0) + program_pointer: uintp = uintp(0) + simulation_pointer: uintp = uintp(0) + data_pointer: uintp = uintp(0) + + # Note that the uintp is manually overriden in code_factory. diff --git a/mcdc/object_/particle.py b/mcdc/object_/particle.py index 76e08fef..3d9fa43f 100644 --- a/mcdc/object_/particle.py +++ b/mcdc/object_/particle.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field from typing import Annotated -from numpy import int64, uint +from numpy import int64, uint64 from numpy.typing import NDArray #### @@ -25,7 +25,7 @@ class ParticleData(ObjectBase): E: float = 0.0 w: float = 0.0 particle_type: int = PARTICLE_NEUTRON - rng_seed: uint = uint(1) + rng_seed: uint64 = uint64(1) @dataclass diff --git a/mcdc/object_/settings.py b/mcdc/object_/settings.py index 460b7de3..60a28d89 100644 --- a/mcdc/object_/settings.py +++ b/mcdc/object_/settings.py @@ -28,10 +28,6 @@ class Settings(ObjectSingleton): N_batch: int = 1 rng_seed: int = 1 - # Simulation mode - multigroup_mode: bool = False - eigenvalue_mode: bool = False - # k-eigenvalue N_inactive: int = 0 N_active: int = 0 @@ -67,6 +63,15 @@ class Settings(ObjectSingleton): electron_transport: bool = False proton_transport: bool = False + # Neutron transport modes + neutron_multigroup_mode: bool = False + neutron_eigenvalue_mode: bool = False + + # GPU mode + gpu_strategy: int = GPU_STRATEGY_ASYNC + gpu_async_type: int = GPU_ASYNC_SIMPLE + gpu_storage: int = GPU_STORAGE_SEPARATE + def __post_init__(self): super().__init__() @@ -104,7 +109,7 @@ def set_eigenmode( self.N_inactive = N_inactive self.N_active = N_active self.N_cycle = self.N_inactive + self.N_active - self.eigenvalue_mode = True + self.neutron_eigenvalue_mode = True self.k_init = k_init self.save_particle = save_particle diff --git a/mcdc/object_/simulation.py b/mcdc/object_/simulation.py index 43e9b704..fa5b5f9b 100644 --- a/mcdc/object_/simulation.py +++ b/mcdc/object_/simulation.py @@ -130,6 +130,7 @@ class Simulation(ObjectSingleton): # GPU metadata gpu_meta: GPUMeta + source_seed: int def __init__(self): super().__init__() @@ -229,6 +230,7 @@ def __init__(self): # GPU metadata self.gpu_meta = GPUMeta() + self.source_seed = 0 def set_root_universe(self, cells=[]): self.universes[0].cells = cells diff --git a/mcdc/output.py b/mcdc/output.py index 5a26a99b..24530c46 100644 --- a/mcdc/output.py +++ b/mcdc/output.py @@ -48,7 +48,7 @@ def generate_output(mcdc, data): create_tally_dataset(file, mcdc, data) # Eigenvalues - if mcdc["settings"]["eigenvalue_mode"]: + if mcdc["settings"]["neutron_eigenvalue_mode"]: N_cycle = mcdc["settings"]["N_cycle"] file.create_dataset( "k_cycle", data=mcdc_get.simulation.k_cycle_chunk(0, N_cycle, mcdc, data) diff --git a/mcdc/print_.py b/mcdc/print_.py index 4da5824e..7a381d0f 100644 --- a/mcdc/print_.py +++ b/mcdc/print_.py @@ -69,8 +69,8 @@ def print_configuration(): sys.stdout.flush() -def print_eigenvalue_header(mcdc): - if mcdc["settings"]["use_gyration_radius"]: +def print_eigenvalue_header(simulation): + if simulation["settings"]["use_gyration_radius"]: print("\n # k GyRad. k (avg) ") print(" ==== ======= ====== ===================") else: @@ -95,11 +95,11 @@ def print_time(tag, t, percent): print(" %s | %.2f seconds (%.1f%%)" % (tag, t, percent)) -def print_runtime(mcdc): - total = mcdc["runtime_total"] - preparation = mcdc["runtime_preparation"] - simulation = mcdc["runtime_simulation"] - output = mcdc["runtime_output"] +def print_runtime(simulation): + total = simulation["runtime_total"] + preparation = simulation["runtime_preparation"] + simulation = simulation["runtime_simulation"] + output = simulation["runtime_output"] print("\n Runtime report:") print_time("Total ", total, 100) print_time("Preparation", preparation, preparation / total * 100) @@ -137,23 +137,23 @@ def print_warning(msg): sys.stdout.flush() -def print_progress(percent, mcdc): +def print_progress(percent, simulation): if master: sys.stdout.write("\r") - if not mcdc["settings"]["eigenvalue_mode"]: - if mcdc["settings"]["N_census"] == 1: + if not simulation["settings"]["neutron_eigenvalue_mode"]: + if simulation["settings"]["N_census"] == 1: sys.stdout.write( " [%-28s] %d%%" % ("=" * int(percent * 28), percent * 100.0) ) else: - idx = mcdc["idx_census"] + 1 - N = mcdc["settings"]["N_census"] + idx = simulation["idx_census"] + 1 + N = simulation["settings"]["N_census"] sys.stdout.write( " Census %i/%i: [%-28s] %d%%" % (idx, N, "=" * int(percent * 28), percent * 100.0) ) else: - if mcdc["settings"]["use_gyration_radius"]: + if simulation["settings"]["use_gyration_radius"]: sys.stdout.write( " [%-40s] %d%%" % ("=" * int(percent * 40), percent * 100.0) ) @@ -164,9 +164,9 @@ def print_progress(percent, mcdc): sys.stdout.flush() -def print_header_eigenvalue(mcdc): +def print_header_eigenvalue(simulation): if master: - if mcdc["settings"]["use_gyration_radius"]: + if simulation["settings"]["use_gyration_radius"]: print("\n # k GyRad. k (avg) ") print(" ==== ======= ====== ===================") else: @@ -180,18 +180,18 @@ def print_header_batch(i, N): sys.stdout.flush() -def print_progress_eigenvalue(mcdc, data): +def print_progress_eigenvalue(simulation, data): if master: - idx_cycle = mcdc["idx_cycle"] - k_eff = mcdc["k_eff"] - k_avg = mcdc["k_avg_running"] - k_sdv = mcdc["k_sdv_running"] - gr = mcdc_get.simulation.gyration_radius(idx_cycle, mcdc, data) - if mcdc["settings"]["use_progress_bar"]: + idx_cycle = simulation["idx_cycle"] + k_eff = simulation["k_eff"] + k_avg = simulation["k_avg_running"] + k_sdv = simulation["k_sdv_running"] + gr = mcdc_get.simulation.gyration_radius(idx_cycle, simulation, data) + if simulation["settings"]["use_progress_bar"]: sys.stdout.write("\r") sys.stdout.write("\033[K") - if mcdc["settings"]["use_gyration_radius"]: - if not mcdc["cycle_active"]: + if simulation["settings"]["use_gyration_radius"]: + if not simulation["cycle_active"]: print(" %-4i %.5f %6.2f" % (idx_cycle + 1, k_eff, gr)) else: print( @@ -199,7 +199,7 @@ def print_progress_eigenvalue(mcdc, data): % (idx_cycle + 1, k_eff, gr, k_avg, k_sdv) ) else: - if not mcdc["cycle_active"]: + if not simulation["cycle_active"]: print(" %-4i %.5f" % (idx_cycle + 1, k_eff)) else: print( @@ -207,17 +207,17 @@ def print_progress_eigenvalue(mcdc, data): ) -def print_runtime(mcdc): - total = mcdc["runtime_total"] - preparation = mcdc["runtime_preparation"] - simulation = mcdc["runtime_simulation"] - output = mcdc["runtime_output"] +def print_runtime(simulation): + t_total = simulation["runtime_total"] + t_preparation = simulation["runtime_preparation"] + t_simulation = simulation["runtime_simulation"] + t_output = simulation["runtime_output"] if master: print("\n Runtime report:") - print_time("Total ", total, 100) - print_time("Preparation", preparation, preparation / total * 100) - print_time("Simulation ", simulation, simulation / total * 100) - print_time("Output ", output, output / total * 100) + print_time("Total ", t_total, 100) + print_time("Preparation", t_preparation, t_preparation / t_total * 100) + print_time("Simulation ", t_simulation, t_simulation / t_total * 100) + print_time("Output ", t_output, t_output / t_total * 100) print("\n") sys.stdout.flush() diff --git a/mcdc/transport/.gitignore b/mcdc/transport/.gitignore deleted file mode 100644 index b9792f1c..00000000 --- a/mcdc/transport/.gitignore +++ /dev/null @@ -1 +0,0 @@ -literals.py diff --git a/mcdc/transport/__init__.py b/mcdc/transport/__init__.py new file mode 100644 index 00000000..c2e7b125 --- /dev/null +++ b/mcdc/transport/__init__.py @@ -0,0 +1,4 @@ +import mcdc.transport.geometry as geometry +import mcdc.transport.particle_bank as particle_bank +import mcdc.transport.simulation as simulation +import mcdc.transport.util as util diff --git a/mcdc/transport/data.py b/mcdc/transport/data.py index 51e8b9b3..443dc629 100644 --- a/mcdc/transport/data.py +++ b/mcdc/transport/data.py @@ -14,14 +14,14 @@ @njit -def evaluate_data(x, data_base, mcdc, data): +def evaluate_data(x, data_base, simulation, data): data_type = data_base["child_type"] ID = data_base["child_ID"] if data_type == DATA_TABLE: - table = mcdc["table_data"][ID] + table = simulation["table_data"][ID] return evaluate_table(x, table, data) elif data_type == DATA_POLYNOMIAL: - polynomial = mcdc["polynomial_data"][ID] + polynomial = simulation["polynomial_data"][ID] return evaluate_polynomial(x, polynomial, data) else: return 0.0 @@ -29,7 +29,11 @@ def evaluate_data(x, data_base, mcdc, data): @njit def evaluate_table(x, table, data): - grid = mcdc_get.table_data.x_all(table, data) + offset = table["x_offset"] + length = table["x_length"] + grid = data[offset : offset + length] + # Above is equivalent to: grid = mcdc_get.table_data.x_all(table, data) + idx = find_bin(x, grid) x1 = grid[idx] x2 = grid[idx + 1] @@ -44,7 +48,11 @@ def evaluate_table(x, table, data): @njit def evaluate_polynomial(x, polynomial, data): - coeffs = mcdc_get.polynomial_data.coefficients_all(polynomial, data) + offset = polynomial["coefficients_offset"] + length = polynomial["coefficients_length"] + coeffs = data[offset : offset + length] + # Above is equivalent to: coeffs = mcdc_get.polynomial_data.coefficients_all(polynomial, data) + total = 0.0 for i in range(len(coeffs)): total += coeffs[i] * x**i diff --git a/mcdc/transport/distribution.py b/mcdc/transport/distribution.py index cbd82941..a2bd93e8 100644 --- a/mcdc/transport/distribution.py +++ b/mcdc/transport/distribution.py @@ -27,29 +27,39 @@ @njit -def sample_distribution(E, distribution, rng_state, mcdc, data, scale=False): +def sample_distribution(E, distribution, rng_state, simulation, data): + return _sample_distribution(E, distribution, rng_state, simulation, data, False) + + +@njit +def sample_distribution_with_scale(E, distribution, rng_state, simulation, data): + return _sample_distribution(E, distribution, rng_state, simulation, data, True) + + +@njit +def _sample_distribution(E, distribution, rng_state, simulation, data, scale): distribution_type = distribution["child_type"] ID = distribution["child_ID"] if distribution_type == DISTRIBUTION_TABULATED: - table = mcdc["tabulated_distributions"][ID] + table = simulation["tabulated_distributions"][ID] return sample_tabulated(table, rng_state, data) elif distribution_type == DISTRIBUTION_MULTITABLE: - multi_table = mcdc["multi_table_distributions"][ID] - return sample_multi_table(E, rng_state, multi_table, data, scale) + multi_table = simulation["multi_table_distributions"][ID] + return _sample_multi_table(E, rng_state, multi_table, data, scale) elif distribution_type == DISTRIBUTION_LEVEL_SCATTERING: - level_scattering = mcdc["level_scattering_distributions"][ID] + level_scattering = simulation["level_scattering_distributions"][ID] return sample_level_scattering(E, level_scattering) elif distribution_type == DISTRIBUTION_EVAPORATION: - evaporation = mcdc["evaporation_distributions"][ID] - return sample_evaporation(E, rng_state, evaporation, mcdc, data) + evaporation = simulation["evaporation_distributions"][ID] + return sample_evaporation(E, rng_state, evaporation, simulation, data) elif distribution_type == DISTRIBUTION_MAXWELLIAN: - maxwellian = mcdc["maxwellian_distributions"][ID] - return sample_maxwellian(E, rng_state, maxwellian, mcdc, data) + maxwellian = simulation["maxwellian_distributions"][ID] + return sample_maxwellian(E, rng_state, maxwellian, simulation, data) # TODO: Should not get here else: @@ -57,20 +67,38 @@ def sample_distribution(E, distribution, rng_state, mcdc, data, scale=False): @njit -def sample_correlated_distribution(E, distribution, rng_state, mcdc, data, scale=False): +def sample_correlated_distribution(E, distribution, rng_state, simulation, data): + return _sample_correlated_distribution( + E, distribution, rng_state, simulation, data, False + ) + + +@njit +def sample_correlated_distribution_with_scale( + E, distribution, rng_state, simulation, data +): + return _sample_correlated_distribution( + E, distribution, rng_state, simulation, data, True + ) + + +@njit +def _sample_correlated_distribution( + E, distribution, rng_state, simulation, data, scale +): distribution_type = distribution["child_type"] ID = distribution["child_ID"] if distribution_type == DISTRIBUTION_KALBACH_MANN: - kalbach_mann = mcdc["kalbach_mann_distributions"][ID] + kalbach_mann = simulation["kalbach_mann_distributions"][ID] return sample_kalbach_mann(E, rng_state, kalbach_mann, data) elif distribution_type == DISTRIBUTION_TABULATED_ENERGY_ANGLE: - table = mcdc["tabulated_energy_angle_distributions"][ID] + table = simulation["tabulated_energy_angle_distributions"][ID] return sample_tabulated_energy_angle(E, rng_state, table, data) elif distribution_type == DISTRIBUTION_N_BODY: - nbody = mcdc["nbody_distributions"][ID] + nbody = simulation["nbody_distributions"][ID] E_out = sample_tabulated(nbody, rng_state, data) mu = sample_isotropic_cosine(rng_state) return E_out, mu @@ -148,7 +176,13 @@ def sample_direction(polar_cosine, azimuthal, polar_coordinate, rng_state): @njit def sample_tabulated(table, rng_state, data): xi = rng.lcg(rng_state) - idx = find_bin(xi, mcdc_get.tabulated_distribution.cdf_all(table, data)) + + offset = table["cdf_offset"] + length = table["cdf_length"] + cdf = data[offset : offset + length] + # Above is equivalent to: cdf = mcdc_get.tabulated_distribution.cdf_all(table, data) + + idx = find_bin(xi, cdf) cdf_low = mcdc_get.tabulated_distribution.cdf(idx, table, data) cdf_high = mcdc_get.tabulated_distribution.cdf(idx + 1, table, data) value_low = mcdc_get.tabulated_distribution.value(idx, table, data) @@ -159,7 +193,13 @@ def sample_tabulated(table, rng_state, data): @njit def sample_pmf(pmf, rng_state, data): xi = rng.lcg(rng_state) - idx = find_bin(xi, mcdc_get.pmf_distribution.cmf_all(pmf, data)) + + offset = pmf["cmf_offset"] + length = pmf["cmf_length"] + cmf = data[offset : offset + length] + # Above is equivalent to: cmf = mcdc_get.pmf_distribution.cmf_all(pmf, data) + + idx = find_bin(xi, cmf) return mcdc_get.pmf_distribution.value(idx, pmf, data) @@ -194,8 +234,16 @@ def sample_white_direction(nx, ny, nz, rng_state): @njit -def sample_multi_table(E, rng_state, multi_table, data, scale=False): - grid = mcdc_get.multi_table_distribution.grid_all(multi_table, data) +def sample_multi_table(E, rng_state, multi_table, data): + return _sample_multi_table(E, rng_state, multi_table, data, False) + + +@njit +def _sample_multi_table(E, rng_state, multi_table, data, scale): + offset = multi_table["grid_offset"] + length = multi_table["grid_length"] + grid = data[offset : offset + length] + # Above is equivalent to: grid = mcdc_get.multi_table_distribution.grid_all(multi_table, data) # Edge cases if E < grid[0]: @@ -257,7 +305,9 @@ def sample_multi_table(E, rng_state, multi_table, data, scale=False): size = end - start # The CDF - cdf = mcdc_get.multi_table_distribution.cdf_chunk(start, size, multi_table, data) + offset = multi_table["cdf_offset"] + cdf = data[start + offset : start + offset + size] + # Above is equivalent to: cdf = mcdc_get.multi_table_distribution.cdf_chunk(start, size, multi_table, data) # Generate random numbers xi = rng.lcg(rng_state) @@ -289,9 +339,9 @@ def sample_multi_table(E, rng_state, multi_table, data, scale=False): @njit -def sample_maxwellian(E, rng_state, maxwellian, mcdc, data): +def sample_maxwellian(E, rng_state, maxwellian, simulation, data): # Get nuclear temperature - table = mcdc["table_data"][maxwellian["nuclear_temperature_ID"]] + table = simulation["table_data"][maxwellian["nuclear_temperature_ID"]] nuclear_temperature = evaluate_table(E, table, data) restriction_energy = maxwellian["restriction_energy"] @@ -319,9 +369,9 @@ def sample_level_scattering(E, level_scattering): @njit -def sample_evaporation(E, rng_state, evaporation, mcdc, data): +def sample_evaporation(E, rng_state, evaporation, simulation, data): # Get nuclear temperature - table = mcdc["table_data"][evaporation["nuclear_temperature_ID"]] + table = simulation["table_data"][evaporation["nuclear_temperature_ID"]] nuclear_temperature = evaluate_table(E, table, data) restriction_energy = evaporation["restriction_energy"] @@ -343,7 +393,10 @@ def sample_evaporation(E, rng_state, evaporation, mcdc, data): @njit def sample_kalbach_mann(E, rng_state, kalbach_mann, data): - grid = mcdc_get.kalbach_mann_distribution.energy_all(kalbach_mann, data) + offset = kalbach_mann["energy_offset"] + length = kalbach_mann["energy_length"] + grid = data[offset : offset + length] + # Above is equivalent to: grid = mcdc_get.kalbach_mann_distribution.energy_all(kalbach_mann, data) # Random numbers xi1 = rng.lcg(rng_state) @@ -397,7 +450,9 @@ def sample_kalbach_mann(E, rng_state, kalbach_mann, data): size = end - start # The CDF - cdf = mcdc_get.kalbach_mann_distribution.cdf_chunk(start, size, kalbach_mann, data) + offset = kalbach_mann["cdf_offset"] + cdf = data[start + offset : start + offset + size] + # Above is equivalent to: cdf = mcdc_get.kalbach_mann_distribution.cdf_chunk(start, size, kalbach_mann, data) # Sample bin index idx = find_bin(xi2, cdf) @@ -446,7 +501,10 @@ def sample_kalbach_mann(E, rng_state, kalbach_mann, data): @njit def sample_tabulated_energy_angle(E, rng_state, table, data): - grid = mcdc_get.tabulated_energy_angle_distribution.energy_all(table, data) + offset = table["energy_offset"] + length = table["energy_length"] + grid = data[offset : offset + length] + # Above is equivalent to: grid = mcdc_get.tabulated_energy_angle_distribution.energy_all(table, data) # Random numbers xi1 = rng.lcg(rng_state) @@ -503,9 +561,12 @@ def sample_tabulated_energy_angle(E, rng_state, table, data): size = end - start # The CDF - cdf = mcdc_get.tabulated_energy_angle_distribution.cdf_chunk( - start, size, table, data - ) + offset = table["cdf_offset"] + cdf = data[start + offset : start + offset + size] + # Above is equivalent to: + # cdf = mcdc_get.tabulated_energy_angle_distribution.cdf_chunk( + # start, size, table, data + # ) # Sample bin index idx = find_bin(xi2, cdf) @@ -555,9 +616,12 @@ def sample_tabulated_energy_angle(E, rng_state, table, data): size = end - start # The CDF - cdf = mcdc_get.tabulated_energy_angle_distribution.cosine_cdf_chunk( - start, size, table, data - ) + offset = table["cosine_cdf_offset"] + cdf = data[start + offset : start + offset + size] + # Above is equivalent to: + # cdf = mcdc_get.tabulated_energy_angle_distribution.cosine_cdf_chunk( + # start, size, table, data + # ) # Sample bin index idx = find_bin(xi3, cdf) diff --git a/mcdc/transport/geometry/__init__.py b/mcdc/transport/geometry/__init__.py index 4153e038..a785177f 100644 --- a/mcdc/transport/geometry/__init__.py +++ b/mcdc/transport/geometry/__init__.py @@ -1,3 +1,5 @@ +import mcdc.transport.geometry.interface as interface + from .interface import ( inspect_geometry, locate_particle, diff --git a/mcdc/transport/geometry/interface.py b/mcdc/transport/geometry/interface.py index fae78e40..0de3f79a 100644 --- a/mcdc/transport/geometry/interface.py +++ b/mcdc/transport/geometry/interface.py @@ -6,10 +6,11 @@ #### import mcdc.mcdc_get as mcdc_get -import mcdc.transport.literals as literals +import mcdc.literals as literals import mcdc.transport.mesh as mesh import mcdc.transport.physics as physics import mcdc.transport.tally as tally_module +import mcdc.transport.util as util from mcdc.constant import * from mcdc.transport.geometry.surface import get_distance, check_sense, reflect @@ -20,7 +21,7 @@ @njit -def inspect_geometry(particle_container, mcdc, data): +def inspect_geometry(particle_container, simulation, data): """ Full geometry inspection of the particle: - Set particle top cell and material IDs (if not lost) @@ -39,7 +40,7 @@ def inspect_geometry(particle_container, mcdc, data): ux_global = particle["ux"] uy_global = particle["uy"] uz_global = particle["uz"] - speed = physics.particle_speed(particle_container, mcdc, data) + speed = physics.particle_speed(particle_container, simulation, data) # Default returns distance = INF @@ -47,20 +48,22 @@ def inspect_geometry(particle_container, mcdc, data): # Find top cell from root universe if unknown if particle["cell_ID"] == -1: - particle["cell_ID"] = get_cell(particle_container, UNIVERSE_ROOT, mcdc, data) + particle["cell_ID"] = get_cell( + particle_container, UNIVERSE_ROOT, simulation, data + ) # Particle is lost? if particle["cell_ID"] == -1: event = EVENT_LOST # The top cell - cell = mcdc["cells"][particle["cell_ID"]] + cell = simulation["cells"][particle["cell_ID"]] # Recursively check cells until material cell is found (or the particle is lost) while event != EVENT_LOST: # Distance to nearest surface d_surface, surface_ID = distance_to_nearest_surface( - particle_container, cell, mcdc, data + particle_container, cell, simulation, data ) # Check if smaller @@ -103,7 +106,7 @@ def inspect_geometry(particle_container, mcdc, data): # Lattice cell? elif cell["fill_type"] == FILL_LATTICE: # Get lattice - lattice = mcdc["lattices"][cell["fill_ID"]] + lattice = simulation["lattices"][cell["fill_ID"]] # Distance to lattice grid d_lattice = mesh.uniform.get_crossing_distance( @@ -137,9 +140,9 @@ def inspect_geometry(particle_container, mcdc, data): particle["z"] -= lattice["z0"] + (iz + 0.5) * lattice["dz"] # Get inner cell - cell_ID = get_cell(particle_container, universe_ID, mcdc, data) + cell_ID = get_cell(particle_container, universe_ID, simulation, data) if cell_ID > -1: - cell = mcdc["cells"][cell_ID] + cell = simulation["cells"][cell_ID] else: event = EVENT_LOST @@ -154,7 +157,7 @@ def inspect_geometry(particle_container, mcdc, data): # Report lost particle if event == EVENT_LOST: - report_lost_particle(particle_container, mcdc) + report_lost_particle(particle_container, simulation) # Assign particle event particle["event"] = event @@ -163,7 +166,7 @@ def inspect_geometry(particle_container, mcdc, data): @njit -def locate_particle(particle_container, mcdc, data): +def locate_particle(particle_container, simulation, data): """ Set particle cell and material IDs Return False if particle is lost @@ -187,14 +190,16 @@ def locate_particle(particle_container, mcdc, data): # Find top cell from root universe if unknown if particle["cell_ID"] == -1: - particle["cell_ID"] = get_cell(particle_container, UNIVERSE_ROOT, mcdc, data) + particle["cell_ID"] = get_cell( + particle_container, UNIVERSE_ROOT, simulation, data + ) # Particle is lost? if particle["cell_ID"] == -1: particle_is_lost = True # The top cell - cell = mcdc["cells"][particle["cell_ID"]] + cell = simulation["cells"][particle["cell_ID"]] # Recursively check cells until material cell is found (or the particle is lost) while not particle_is_lost: @@ -224,7 +229,7 @@ def locate_particle(particle_container, mcdc, data): # Lattice cell? elif cell["fill_type"] == FILL_LATTICE: # Get lattice - lattice = mcdc["lattices"][cell["fill_ID"]] + lattice = simulation["lattices"][cell["fill_ID"]] # Get universe ix, iy, iz = mesh.uniform.get_indices(particle_container, lattice) @@ -241,9 +246,9 @@ def locate_particle(particle_container, mcdc, data): particle["z"] -= lattice["z0"] + (iz + 0.5) * lattice["dz"] # Get inner cell - cell_ID = get_cell(particle_container, universe_ID, mcdc, data) + cell_ID = get_cell(particle_container, universe_ID, simulation, data) if cell_ID > -1: - cell = mcdc["cells"][cell_ID] + cell = simulation["cells"][cell_ID] else: particle_is_lost = True @@ -258,7 +263,7 @@ def locate_particle(particle_container, mcdc, data): # Report lost particle if particle_is_lost: - report_lost_particle(particle_container, mcdc) + report_lost_particle(particle_container, simulation) return not particle_is_lost @@ -325,19 +330,19 @@ def _rotation_matrix(rotation): @njit -def get_cell(particle_container, universe_ID, mcdc, data): +def get_cell(particle_container, universe_ID, simulation, data): """ Find and return particle cell ID in the given universe Return -1 if particle is lost """ particle = particle_container[0] - universe = mcdc["universes"][universe_ID] + universe = simulation["universes"][universe_ID] # Check over all cells in the universe for i in range(universe["N_cell"]): cell_ID = int(mcdc_get.universe.cell_IDs(i, universe, data)) - cell = mcdc["cells"][cell_ID] - if check_cell(particle_container, cell, mcdc, data): + cell = simulation["cells"][cell_ID] + if check_cell(particle_container, cell, simulation, data): return cell_ID # Particle is not found @@ -345,7 +350,7 @@ def get_cell(particle_container, universe_ID, mcdc, data): @njit -def check_cell(particle_container, cell, mcdc, data): +def check_cell(particle_container, cell, simulation, data): """ Check if the particle is inside the cell """ @@ -357,18 +362,18 @@ def check_cell(particle_container, cell, mcdc, data): return True # Create local value array - value = np.zeros(literals.rpn_evaluation_buffer_size, np.bool_) + value = util.local_array(literals.rpn_evaluation_buffer_size(), np.bool_) N_value = 0 # Particle parameters - speed = physics.particle_speed(particle_container, mcdc, data) + speed = physics.particle_speed(particle_container, simulation, data) # March forward through RPN tokens for idx in range(N_token): token = int(mcdc_get.cell.region_RPN_tokens(idx, cell, data)) if token >= 0: - surface = mcdc["surfaces"][token] + surface = simulation["surfaces"][token] value[N_value] = check_sense(particle_container, speed, surface, data) N_value += 1 @@ -387,7 +392,7 @@ def check_cell(particle_container, cell, mcdc, data): @njit -def report_lost_particle(particle_container, mcdc): +def report_lost_particle(particle_container, simulation): """ Report lost particle and terminate it """ @@ -397,9 +402,9 @@ def report_lost_particle(particle_container, mcdc): y = particle["y"] z = particle["z"] t = particle["t"] - idx_batch = mcdc["idx_batch"] - idx_census = mcdc["idx_census"] - idx_work = mcdc["idx_work"] + idx_batch = simulation["idx_batch"] + idx_census = simulation["idx_census"] + idx_work = simulation["idx_work"] print("A particle is lost at (", x, y, z, t, ")") print(" (batch/census/work) indices: (", idx_batch, idx_census, idx_work, ")") particle["alive"] = False @@ -411,7 +416,7 @@ def report_lost_particle(particle_container, mcdc): @njit -def distance_to_nearest_surface(particle_container, cell, mcdc, data): +def distance_to_nearest_surface(particle_container, cell, simulation, data): """ Determine the nearest cell surface and the distance to it """ @@ -419,12 +424,12 @@ def distance_to_nearest_surface(particle_container, cell, mcdc, data): surface_ID = -1 # Particle parameters - speed = physics.particle_speed(particle_container, mcdc, data) + speed = physics.particle_speed(particle_container, simulation, data) # Iterate over all surfaces and find the minimum distance for i in range(cell["N_surface"]): candidate_surface_ID = int(mcdc_get.cell.surface_IDs(i, cell, data)) - surface = mcdc["surfaces"][candidate_surface_ID] + surface = simulation["surfaces"][candidate_surface_ID] d = get_distance(particle_container, speed, surface, data) if d < distance: distance = d @@ -434,11 +439,11 @@ def distance_to_nearest_surface(particle_container, cell, mcdc, data): @njit -def surface_crossing(P_arr, mcdc, data): +def surface_crossing(P_arr, simulation, data): P = P_arr[0] # Apply BC - surface = mcdc["surfaces"][P["surface_ID"]] + surface = simulation["surfaces"][P["surface_ID"]] BC = surface["boundary_condition"] if BC == BC_VACUUM: P["alive"] = False @@ -448,8 +453,8 @@ def surface_crossing(P_arr, mcdc, data): # Score tally for i in range(surface["N_tally"]): tally_ID = int(mcdc_get.surface.tally_IDs(i, surface, data)) - tally = mcdc["surface_tallies"][tally_ID] - tally_module.score.surface_tally(P_arr, surface, tally, mcdc, data) + tally = simulation["surface_tallies"][tally_ID] + tally_module.score.surface_tally(P_arr, surface, tally, simulation, data) # Need to check new cell later? if P["alive"] and not BC == BC_REFLECTIVE: diff --git a/mcdc/transport/geometry/surface/interface.py b/mcdc/transport/geometry/surface/interface.py index 33392d6d..e3adf14d 100644 --- a/mcdc/transport/geometry/surface/interface.py +++ b/mcdc/transport/geometry/surface/interface.py @@ -38,7 +38,7 @@ SURFACE_CONE_Z, SURFACE_TORUS_Z, ) -from mcdc.transport.util import find_bin +from mcdc.transport.util import find_bin_with_rules @njit @@ -360,8 +360,15 @@ def _get_move_idx(t, surface, data): """ Get moving interval index wrt the given time """ - time_grid = mcdc_get.surface.move_time_grid_all(surface, data) - idx = find_bin(t, time_grid, epsilon=COINCIDENCE_TOLERANCE_TIME, go_lower=False) + time_grid = data[ + surface["move_time_grid_offset"] : ( + surface["move_time_grid_offset"] + surface["N_move_grid"] + ) + ] + # Above is equivalent to: time_grid = mcdc_get.surface.move_time_grid_all(surface, data) + tolerance = COINCIDENCE_TOLERANCE_TIME + go_lower = False + idx = find_bin_with_rules(t, time_grid, tolerance, go_lower) # Coinciding cases if abs(time_grid[idx + 1] - t) < COINCIDENCE_TOLERANCE: @@ -377,10 +384,18 @@ def _translate_particle_position(particle_container, surface, idx, data): """ particle = particle_container[0] - # Surface move translations, velocities, and time grid - trans_0 = mcdc_get.surface.move_translations_vector(idx, surface, data) + # Surface move translations + start = surface["move_translations_offset"] + idx * 3 + trans_0 = data[start : start + 3] + # Above is equivalent to: trans_0 = mcdc_get.surface.move_translations_vector(idx, surface, data) + + # Surface move velocities + start = surface["move_velocities_offset"] + idx * 3 + V = data[start : start + 3] + # Above is equivalent to: V = mcdc_get.surface.move_velocities_vector(idx, surface, data) + + # Surface move time grid time_0 = mcdc_get.surface.move_time_grid(idx, surface, data) - V = mcdc_get.surface.move_velocities_vector(idx, surface, data) # Translate the particle t_local = particle["t"] - time_0 @@ -396,8 +411,10 @@ def _translate_particle_direction(particle_container, speed, surface, idx, data) """ particle = particle_container[0] - # Surface move translations, velocities, and time grid - V = mcdc_get.surface.move_velocities_vector(idx, surface, data) + # Surface move velocities + start = surface["move_velocities_offset"] + idx * 3 + V = data[start : start + 3] + # Above is equivalent to: V = mcdc_get.surface.move_velocities_vector(idx, surface, data) # Translate the particle particle["ux"] -= V[0] / speed diff --git a/mcdc/transport/literals.py b/mcdc/transport/literals.py new file mode 100644 index 00000000..bfca9dd7 --- /dev/null +++ b/mcdc/transport/literals.py @@ -0,0 +1,3 @@ +# The following is automatically generated by code_factory.py + +rpn_evaluation_buffer_size = 2 diff --git a/mcdc/transport/mesh/interface.py b/mcdc/transport/mesh/interface.py index 97cdc416..2e4fda69 100644 --- a/mcdc/transport/mesh/interface.py +++ b/mcdc/transport/mesh/interface.py @@ -7,57 +7,57 @@ @njit -def get_indices(particle_container, mesh_base, mcdc, data): +def get_indices(particle_container, mesh_base, simulation, data): mesh_type = mesh_base["child_type"] mesh_ID = mesh_base["child_ID"] if mesh_type == MESH_UNIFORM: - mesh = mcdc["uniform_meshes"][mesh_ID] + mesh = simulation["uniform_meshes"][mesh_ID] return uniform.get_indices(particle_container, mesh) elif mesh_type == MESH_STRUCTURED: - mesh = mcdc["structured_meshes"][mesh_ID] + mesh = simulation["structured_meshes"][mesh_ID] return structured.get_indices(particle_container, mesh, data) return -1, -1, -1 @njit -def get_x(index, mesh_base, mcdc, data): +def get_x(index, mesh_base, simulation, data): mesh_type = mesh_base["child_type"] mesh_ID = mesh_base["child_ID"] if mesh_type == MESH_UNIFORM: - mesh = mcdc["uniform_meshes"][mesh_ID] + mesh = simulation["uniform_meshes"][mesh_ID] return mesh["x0"] + mesh["dx"] * index elif mesh_type == MESH_STRUCTURED: - mesh = mcdc["structured_meshes"][mesh_ID] + mesh = simulation["structured_meshes"][mesh_ID] return mcdc_get.structured_mesh.x(index, mesh, data) return 0.0 @njit -def get_y(index, mesh_base, mcdc, data): +def get_y(index, mesh_base, simulation, data): mesh_type = mesh_base["child_type"] mesh_ID = mesh_base["child_ID"] if mesh_type == MESH_UNIFORM: - mesh = mcdc["uniform_meshes"][mesh_ID] + mesh = simulation["uniform_meshes"][mesh_ID] return mesh["y0"] + mesh["dy"] * index elif mesh_type == MESH_STRUCTURED: - mesh = mcdc["structured_meshes"][mesh_ID] + mesh = simulation["structured_meshes"][mesh_ID] return mcdc_get.structured_mesh.y(index, mesh, data) return 0.0 @njit -def get_z(index, mesh_base, mcdc, data): +def get_z(index, mesh_base, simulation, data): mesh_type = mesh_base["child_type"] mesh_ID = mesh_base["child_ID"] if mesh_type == MESH_UNIFORM: - mesh = mcdc["uniform_meshes"][mesh_ID] + mesh = simulation["uniform_meshes"][mesh_ID] return mesh["z0"] + mesh["dz"] * index elif mesh_type == MESH_STRUCTURED: - mesh = mcdc["structured_meshes"][mesh_ID] + mesh = simulation["structured_meshes"][mesh_ID] return mcdc_get.structured_mesh.z(index, mesh, data) return 0.0 diff --git a/mcdc/transport/mesh/structured.py b/mcdc/transport/mesh/structured.py index 7ea77863..9afe50a8 100644 --- a/mcdc/transport/mesh/structured.py +++ b/mcdc/transport/mesh/structured.py @@ -5,7 +5,7 @@ import mcdc.mcdc_get as mcdc_get from mcdc.constant import COINCIDENCE_TOLERANCE, COINCIDENCE_TOLERANCE_TIME, INF -from mcdc.transport.util import find_bin +from mcdc.transport.util import find_bin_with_rules @njit @@ -23,10 +23,21 @@ def get_indices(particle_container, mesh, data): uy = particle["uy"] uz = particle["uz"] + grid_x = data[mesh["x_offset"] : (mesh["x_offset"] + mesh["x_length"])] + # Above is equivalent to: grid_x = mcdc_get.structured_mesh.x_all(mesh, data) + grid_y = data[mesh["y_offset"] : (mesh["y_offset"] + mesh["y_length"])] + # Above is equivalent to: grid_y = mcdc_get.structured_mesh.y_all(mesh, data) + grid_z = data[mesh["z_offset"] : (mesh["z_offset"] + mesh["z_length"])] + # Above is equivalent to: grid_z = mcdc_get.structured_mesh.z_all(mesh, data) + tolerance = COINCIDENCE_TOLERANCE - ix = find_bin(x, mcdc_get.structured_mesh.x_all(mesh, data), tolerance, ux < 0.0) - iy = find_bin(y, mcdc_get.structured_mesh.y_all(mesh, data), tolerance, uy < 0.0) - iz = find_bin(z, mcdc_get.structured_mesh.z_all(mesh, data), tolerance, uz < 0.0) + ux_go_lower = ux < 0.0 + uy_go_lower = uy < 0.0 + uz_go_lower = uz < 0.0 + + ix = find_bin_with_rules(x, grid_x, tolerance, ux_go_lower) + iy = find_bin_with_rules(y, grid_y, tolerance, uy_go_lower) + iz = find_bin_with_rules(z, grid_z, tolerance, uz_go_lower) return ix, iy, iz diff --git a/mcdc/transport/mpi.py b/mcdc/transport/mpi.py index 395957f6..35b38c41 100644 --- a/mcdc/transport/mpi.py +++ b/mcdc/transport/mpi.py @@ -4,9 +4,9 @@ @njit -def distribute_work(N_work, mcdc): - size = mcdc["mpi_size"] - rank = mcdc["mpi_rank"] +def distribute_work(N_work, simulation): + size = simulation["mpi_size"] + rank = simulation["mpi_rank"] # Total number of work work_size_total = N_work @@ -28,6 +28,6 @@ def distribute_work(N_work, mcdc): work_start += rem # Store the workload specification - mcdc["mpi_work_start"] = work_start - mcdc["mpi_work_size"] = work_size - mcdc["mpi_work_size_total"] = work_size_total + simulation["mpi_work_start"] = work_start + simulation["mpi_work_size"] = work_size + simulation["mpi_work_size_total"] = work_size_total diff --git a/mcdc/transport/particle.py b/mcdc/transport/particle.py index 5a01b02b..e524e1c2 100644 --- a/mcdc/transport/particle.py +++ b/mcdc/transport/particle.py @@ -7,9 +7,9 @@ @njit -def move(particle_container, distance, mcdc, data): +def move(particle_container, distance, simulation, data): particle = particle_container[0] - ut = 1.0 / physics.particle_speed(particle_container, mcdc, data) + ut = 1.0 / physics.particle_speed(particle_container, simulation, data) particle["x"] += particle["ux"] * distance particle["y"] += particle["uy"] * distance diff --git a/mcdc/transport/particle_bank.py b/mcdc/transport/particle_bank.py index 13f0b550..9b32f9cd 100644 --- a/mcdc/transport/particle_bank.py +++ b/mcdc/transport/particle_bank.py @@ -13,10 +13,10 @@ import mcdc.transport.mpi as mpi import mcdc.transport.particle as particle_module import mcdc.transport.technique as technique +import mcdc.transport.util as util from mcdc.constant import * from mcdc.print_ import print_error -from mcdc.transport.util import atomic_add # ============================================================================= # Bank size @@ -35,7 +35,7 @@ def set_bank_size(bank, value): @njit def add_bank_size(bank, value): - atomic_add(bank["size"], 0, value) + util.atomic_add(bank["size"], 0, value) # ============================================================================= @@ -46,35 +46,53 @@ def add_bank_size(bank, value): @njit def _bank_particle(particle_container, bank): # Check if bank is full - if get_bank_size(bank) == bank["particles"].shape[0]: + if get_bank_size(bank) == bank["particle_data"].shape[0]: report_full_bank(bank) # Set particle data idx = get_bank_size(bank) - particle_module.copy(bank["particles"][idx : idx + 1], particle_container) + particle_module.copy(bank["particle_data"][idx : idx + 1], particle_container) + + +@njit +def bank_active_particle(particle_container, program): + simulation = util.access_simulation(program) + bank = simulation["bank_active"] + _bank_particle(particle_container, bank) # Increment bank size add_bank_size(bank, 1) @njit -def bank_active_particle(particle_container, mcdc): - _bank_particle(particle_container, mcdc["bank_active"]) - +def bank_census_particle(particle_container, program): + simulation = util.access_simulation(program) + bank = simulation["bank_census"] + _bank_particle(particle_container, bank) -@njit -def bank_source_particle(particle_container, mcdc): - _bank_particle(particle_container, mcdc["bank_source"]) + # Increment bank size + add_bank_size(bank, 1) @njit -def bank_census_particle(particle_container, mcdc): - _bank_particle(particle_container, mcdc["bank_census"]) +def bank_future_particle(particle_container, program): + simulation = util.access_simulation(program) + bank = simulation["bank_future"] + _bank_particle(particle_container, bank) + + # Increment bank size + add_bank_size(bank, 1) @njit -def bank_future_particle(particle_container, mcdc): - _bank_particle(particle_container, mcdc["bank_future"]) +def bank_source_particle(particle_container, simulation): + bank = simulation["bank_source"] + _bank_particle(particle_container, bank) + + # Increment bank size + # Note that we don't use the atomic operation in add_bank_size function + # as source particle banking is not thread-parallelized + bank["size"][0] += 1 @njit @@ -85,7 +103,7 @@ def pop_particle(particle_container, bank): # Set particle data idx = get_bank_size(bank) - 1 - particle_module.copy(particle_container, bank["particles"][idx : idx + 1]) + particle_module.copy(particle_container, bank["particle_data"][idx : idx + 1]) # Decrement bank size add_bank_size(bank, -1) @@ -117,16 +135,18 @@ def report_empty_bank(bank): @njit -def promote_future_particles(mcdc, data): +def promote_future_particles(program, data): + simulation = util.access_simulation(program) + # Get the banks - future_bank = mcdc["bank_future"] + future_bank = simulation["bank_future"] # Get the next census time - idx = mcdc["idx_census"] + 1 - next_census_time = mcdc_get.settings.census_time(idx, mcdc["settings"], data) + idx = simulation["idx_census"] + 1 + next_census_time = mcdc_get.settings.census_time(idx, simulation["settings"], data) # Particle container - particle_container = np.zeros(1, type_.particle_data) + particle_container = util.local_array(1, type_.particle_data) particle = particle_container[0] # Loop over all particles in future bank @@ -136,19 +156,20 @@ def promote_future_particles(mcdc, data): # NOTE: future bank size decreases as particles are promoted to census bank idx = i - (initial_size - get_bank_size(future_bank)) particle_module.copy( - particle_container, future_bank["particles"][idx : idx + 1] + particle_container, future_bank["particle_data"][idx : idx + 1] ) # Promote the future particle to census bank if particle["t"] < next_census_time: - bank_census_particle(particle_container, mcdc) + + bank_census_particle(particle_container, program) add_bank_size(future_bank, -1) # Consolidate the emptied space in the future bank j = get_bank_size(future_bank) particle_module.copy( - future_bank["particles"][idx : idx + 1], - future_bank["particles"][j : j + 1], + future_bank["particle_data"][idx : idx + 1], + future_bank["particle_data"][j : j + 1], ) @@ -158,52 +179,59 @@ def promote_future_particles(mcdc, data): @njit -def manage_particle_banks(mcdc): - master = mcdc["mpi_master"] - serial = mcdc["mpi_size"] == 1 +def manage_particle_banks(simulation): + master = simulation["mpi_master"] + serial = simulation["mpi_size"] == 1 # TIMER: bank management + time_start = 0.0 if master: with objmode(time_start="float64"): time_start = MPI.Wtime() + time_spent = -time_start # Reset source bank - set_bank_size(mcdc["bank_source"], 0) + set_bank_size(simulation["bank_source"], 0) # Normalize weight - if mcdc["settings"]["eigenvalue_mode"]: - normalize_weight(mcdc["bank_census"], mcdc["settings"]["N_particle"]) + if simulation["settings"]["neutron_eigenvalue_mode"]: + normalize_weight( + simulation["bank_census"], simulation["settings"]["N_particle"] + ) # Population control - if mcdc["population_control"]["active"]: - technique.population_control(mcdc) + if simulation["population_control"]["active"]: + technique.population_control(simulation) else: # Swap census and source bank - source_bank = mcdc["bank_source"] - census_bank = mcdc["bank_census"] + source_bank = simulation["bank_source"] + census_bank = simulation["bank_census"] size = get_bank_size(census_bank) - if size >= source_bank["particles"].shape[0]: + if size >= source_bank["particle_data"].shape[0]: report_full_bank(source_bank) # TODO: better alternative? - source_bank["particles"][:size] = census_bank["particles"][:size] + source_bank["particle_data"][:size] = census_bank["particle_data"][:size] set_bank_size(source_bank, size) # Redistribute work and rebalance bank size across MPI ranks if serial: - mpi.distribute_work(get_bank_size(mcdc["bank_source"]), mcdc) + mpi.distribute_work(get_bank_size(simulation["bank_source"]), simulation) else: - bank_rebalance(mcdc) + bank_rebalance(simulation) # Reset census bank - set_bank_size(mcdc["bank_census"], 0) + set_bank_size(simulation["bank_census"], 0) # TIMER: bank management + time_end = 0.0 if master: with objmode(time_end="float64"): time_end = MPI.Wtime() - mcdc["runtime_bank_management"] += time_end - time_start + time_spent += time_end + if master: + simulation["runtime_bank_management"] += time_spent # ====================================================================================== @@ -212,25 +240,25 @@ def manage_particle_banks(mcdc): @njit -def bank_rebalance(mcdc): +def bank_rebalance(simulation): # Scan the bank - idx_start, N_local, N = bank_scanning(mcdc["bank_source"], mcdc) + idx_start, N_local, N = bank_scanning(simulation["bank_source"], simulation) idx_end = idx_start + N_local - mpi.distribute_work(N, mcdc) + mpi.distribute_work(N, simulation) # Abort if source bank is empty if N == 0: return # Rebalance not needed if there is only one rank - if mcdc["mpi_size"] <= 1: + if simulation["mpi_size"] <= 1: return # Some constants - work_start = mcdc["mpi_work_start"] - work_end = work_start + mcdc["mpi_work_size"] - left = mcdc["mpi_rank"] - 1 - right = mcdc["mpi_rank"] + 1 + work_start = simulation["mpi_work_start"] + work_end = work_start + simulation["mpi_work_size"] + left = simulation["mpi_rank"] - 1 + right = simulation["mpi_rank"] + 1 # Flags if need to receive from or sent to the neighbors send_to_left = idx_start < work_start @@ -247,13 +275,13 @@ def bank_rebalance(mcdc): # MPI nearest-neighbor send/receive buff = np.zeros( - mcdc["bank_source"]["particles"].shape[0], dtype=type_.particle_data + simulation["bank_source"]["particle_data"].shape[0], dtype=type_.particle_data ) with objmode(size="int64"): # Create MPI-supported numpy object - size = get_bank_size(mcdc["bank_source"]) - bank = np.array(mcdc["bank_source"]["particles"][:size]) + size = get_bank_size(simulation["bank_source"]) + bank = np.array(simulation["bank_source"]["particle_data"][:size]) if receive_first: if receive_from_left: @@ -289,9 +317,9 @@ def bank_rebalance(mcdc): buff[i] = bank[i] # Set source bank from buffer - set_bank_size(mcdc["bank_source"], size) + set_bank_size(simulation["bank_source"], size) for i in range(size): - mcdc["bank_source"]["particles"][i] = buff[i] + simulation["bank_source"]["particle_data"][i] = buff[i] # ====================================================================================== @@ -300,7 +328,7 @@ def bank_rebalance(mcdc): @njit -def bank_scanning(bank, mcdc): +def bank_scanning(bank, simulation): N_local = get_bank_size(bank) # Starting index @@ -312,19 +340,19 @@ def bank_scanning(bank, mcdc): # Global size buff[0] += N_local with objmode(): - MPI.COMM_WORLD.Bcast(buff, mcdc["mpi_size"] - 1) + MPI.COMM_WORLD.Bcast(buff, simulation["mpi_size"] - 1) N_global = buff[0] return idx_start, N_local, N_global @njit -def bank_scanning_weight(bank, mcdc): +def bank_scanning_weight(bank, simulation): # Local weight CDF N_local = get_bank_size(bank) w_cdf = np.zeros(N_local + 1) for i in range(N_local): - w_cdf[i + 1] = w_cdf[i] + bank["particles"][i]["w"] + w_cdf[i + 1] = w_cdf[i] + bank["particle_data"][i]["w"] W_local = w_cdf[-1] # Starting weight @@ -337,7 +365,7 @@ def bank_scanning_weight(bank, mcdc): # Global weight buff[0] = w_cdf[-1] with objmode(): - MPI.COMM_WORLD.Bcast(buff, mcdc["mpi_size"] - 1) + MPI.COMM_WORLD.Bcast(buff, simulation["mpi_size"] - 1) W_global = buff[0] return w_start, w_cdf, W_global @@ -350,7 +378,7 @@ def normalize_weight(bank, norm): # Normalize weight for i in range(get_bank_size(bank)): - bank["particles"][i]["w"] *= norm / W + bank["particle_data"][i]["w"] *= norm / W @njit @@ -358,7 +386,7 @@ def total_weight(bank): # Local total weight W_local = np.zeros(1) for i in range(get_bank_size(bank)): - W_local[0] += bank["particles"][i]["w"] + W_local[0] += bank["particle_data"][i]["w"] # MPI Allreduce buff = np.zeros(1, np.float64) diff --git a/mcdc/transport/physics/electron/interface.py b/mcdc/transport/physics/electron/interface.py index 83f6ddde..d94a9759 100644 --- a/mcdc/transport/physics/electron/interface.py +++ b/mcdc/transport/physics/electron/interface.py @@ -10,7 +10,7 @@ @njit -def particle_speed(particle_container, mcdc, data): +def particle_speed(particle_container, simulation, data): return native.particle_speed(particle_container) @@ -20,8 +20,8 @@ def particle_speed(particle_container, mcdc, data): @njit -def macro_xs(reaction_type, particle_container, mcdc, data): - return native.macro_xs(reaction_type, particle_container, mcdc, data) +def macro_xs(reaction_type, particle_container, simulation, data): + return native.macro_xs(reaction_type, particle_container, simulation, data) # ====================================================================================== @@ -30,5 +30,5 @@ def macro_xs(reaction_type, particle_container, mcdc, data): @njit -def collision(particle_container, collision_data_container, mcdc, data): - native.collision(particle_container, collision_data_container, mcdc, data) +def collision(particle_container, collision_data_container, program, data): + native.collision(particle_container, collision_data_container, program, data) diff --git a/mcdc/transport/physics/electron/native.py b/mcdc/transport/physics/electron/native.py index 1e4cba90..4c296ceb 100644 --- a/mcdc/transport/physics/electron/native.py +++ b/mcdc/transport/physics/electron/native.py @@ -10,6 +10,7 @@ import mcdc.transport.particle as particle_module import mcdc.transport.particle_bank as particle_bank_module import mcdc.transport.rng as rng +import mcdc.transport.util as util from mcdc.constant import ( ELECTRON_CUTOFF_ENERGY, @@ -23,7 +24,10 @@ PI, ) from mcdc.transport.data import evaluate_data -from mcdc.transport.distribution import sample_distribution, sample_multi_table +from mcdc.transport.distribution import ( + sample_distribution, + sample_multi_table, +) from mcdc.transport.physics.util import ( evaluate_electron_xs_energy_grid, scatter_direction, @@ -49,15 +53,15 @@ def particle_speed(particle_container): @njit -def macro_xs(reaction_type, particle_container, mcdc, data): +def macro_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] - material = mcdc["native_materials"][particle["material_ID"]] + material = simulation["native_materials"][particle["material_ID"]] E = particle["E"] total = 0.0 for i in range(material["N_element"]): element_ID = int(mcdc_get.native_material.element_IDs(i, material, data)) - element = mcdc["elements"][element_ID] + element = simulation["elements"][element_ID] element_density = mcdc_get.native_material.element_densities(i, material, data) xs = total_micro_xs(reaction_type, E, element, data) @@ -109,10 +113,11 @@ def reaction_micro_xs(E, reaction_base, element, data): @njit -def collision(particle_container, collision_data_container, mcdc, data): +def collision(particle_container, collision_data_container, program, data): + simulation = util.access_simulation(program) particle = particle_container[0] collision_data = collision_data_container[0] - material = mcdc["native_materials"][particle["material_ID"]] + material = simulation["native_materials"][particle["material_ID"]] # Particle properties E = particle["E"] @@ -128,13 +133,13 @@ def collision(particle_container, collision_data_container, mcdc, data): # Sample colliding element # ================================================================================== - SigmaT = macro_xs(ELECTRON_REACTION_TOTAL, particle_container, mcdc, data) + SigmaT = macro_xs(ELECTRON_REACTION_TOTAL, particle_container, simulation, data) xi = rng.lcg(particle_container) * SigmaT total = 0.0 for i in range(material["N_element"]): element_ID = int(mcdc_get.native_material.element_IDs(i, material, data)) - element = mcdc["elements"][element_ID] + element = simulation["elements"][element_ID] element_density = mcdc_get.native_material.element_densities(i, material, data) sigmaT = total_micro_xs(ELECTRON_REACTION_TOTAL, E, element, data) @@ -168,9 +173,9 @@ def collision(particle_container, collision_data_container, mcdc, data): reaction_ID = int( mcdc_get.element.electron_ionization_reaction_IDs(i, element, data) ) - reaction = mcdc["electron_ionization_reactions"][reaction_ID] + reaction = simulation["electron_ionization_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["electron_reactions"][reaction_base_ID] + reaction_base = simulation["electron_reactions"][reaction_base_ID] total += reaction_micro_xs(E, reaction_base, element, data) if xi < total: @@ -179,7 +184,7 @@ def collision(particle_container, collision_data_container, mcdc, data): particle_container, collision_data_container, element, - mcdc, + program, data, ) return @@ -194,13 +199,15 @@ def collision(particle_container, collision_data_container, mcdc, data): i, element, data ) ) - reaction = mcdc["electron_elastic_scattering_reactions"][reaction_ID] + reaction = simulation["electron_elastic_scattering_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["electron_reactions"][reaction_base_ID] + reaction_base = simulation["electron_reactions"][reaction_base_ID] total += reaction_micro_xs(E, reaction_base, element, data) if xi < total: - elastic_scattering(reaction, particle_container, element, mcdc, data) + elastic_scattering( + reaction, particle_container, element, simulation, data + ) return # Bremsstrahlung @@ -211,14 +218,18 @@ def collision(particle_container, collision_data_container, mcdc, data): reaction_ID = int( mcdc_get.element.electron_bremsstrahlung_reaction_IDs(i, element, data) ) - reaction = mcdc["electron_bremsstrahlung_reactions"][reaction_ID] + reaction = simulation["electron_bremsstrahlung_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["electron_reactions"][reaction_base_ID] + reaction_base = simulation["electron_reactions"][reaction_base_ID] total += reaction_micro_xs(E, reaction_base, element, data) if xi < total: bremsstrahlung( - reaction, particle_container, collision_data_container, mcdc, data + reaction, + particle_container, + collision_data_container, + simulation, + data, ) return @@ -230,14 +241,18 @@ def collision(particle_container, collision_data_container, mcdc, data): reaction_ID = int( mcdc_get.element.electron_excitation_reaction_IDs(i, element, data) ) - reaction = mcdc["electron_excitation_reactions"][reaction_ID] + reaction = simulation["electron_excitation_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["electron_reactions"][reaction_base_ID] + reaction_base = simulation["electron_reactions"][reaction_base_ID] total += reaction_micro_xs(E, reaction_base, element, data) if xi < total: excitation( - reaction, particle_container, collision_data_container, mcdc, data + reaction, + particle_container, + collision_data_container, + simulation, + data, ) return @@ -248,7 +263,7 @@ def collision(particle_container, collision_data_container, mcdc, data): @njit -def elastic_scattering(reaction, particle_container, element, mcdc, data): +def elastic_scattering(reaction, particle_container, element, simulation, data): particle = particle_container[0] # Current energy @@ -259,11 +274,11 @@ def elastic_scattering(reaction, particle_container, element, mcdc, data): # ------------------------------------------------------------------------- reaction_base_ID = int(reaction["parent_ID"]) - reaction_base = mcdc["electron_reactions"][reaction_base_ID] + reaction_base = simulation["electron_reactions"][reaction_base_ID] xs_total = reaction_micro_xs(E, reaction_base, element, data) # If large-angle, xs from data table - xs_large = elastic_large_xs(E, reaction, mcdc, data) + xs_large = elastic_large_xs(E, reaction, simulation, data) # Important to check because of numerical issues if xs_large < 0.0: @@ -281,7 +296,7 @@ def elastic_scattering(reaction, particle_container, element, mcdc, data): # Large-angle elastic scattering # --------------------------------------------------------------------- - multi_table = mcdc["multi_table_distributions"][reaction["mu_ID"]] + multi_table = simulation["multi_table_distributions"][reaction["mu_ID"]] mu0 = sample_multi_table(E, particle_container, multi_table, data) else: @@ -334,9 +349,9 @@ def sample_small_angle_mu_coulomb(E, Z, rng_state, mu_cut): @njit -def elastic_large_xs(E, reaction, mcdc, data): - data_base = mcdc["data"][int(reaction["xs_large_ID"])] - return evaluate_data(E, data_base, mcdc, data) +def elastic_large_xs(E, reaction, simulation, data): + data_base = simulation["data"][int(reaction["xs_large_ID"])] + return evaluate_data(E, data_base, simulation, data) # ====================================================================================== @@ -345,14 +360,16 @@ def elastic_large_xs(E, reaction, mcdc, data): @njit -def excitation(reaction, particle_container, collision_data_container, mcdc, data): +def excitation( + reaction, particle_container, collision_data_container, simulation, data +): particle = particle_container[0] collision_data = collision_data_container[0] # Current energy E = particle["E"] - dE = evaluate_eloss(E, reaction, mcdc, data) + dE = evaluate_eloss(E, reaction, simulation, data) # Calculate outgoing energy E_out = E - dE @@ -370,9 +387,9 @@ def excitation(reaction, particle_container, collision_data_container, mcdc, dat @njit -def evaluate_eloss(E, reaction, mcdc, data): - data_base = mcdc["data"][int(reaction["eloss_ID"])] - return evaluate_data(E, data_base, mcdc, data) +def evaluate_eloss(E, reaction, simulation, data): + data_base = simulation["data"][int(reaction["eloss_ID"])] + return evaluate_data(E, data_base, simulation, data) # ====================================================================================== @@ -381,14 +398,16 @@ def evaluate_eloss(E, reaction, mcdc, data): @njit -def bremsstrahlung(reaction, particle_container, collision_data_container, mcdc, data): +def bremsstrahlung( + reaction, particle_container, collision_data_container, simulation, data +): particle = particle_container[0] collision_data = collision_data_container[0] # Current energy E = particle["E"] - dE = evaluate_eloss(E, reaction, mcdc, data) + dE = evaluate_eloss(E, reaction, simulation, data) E_out = E - dE # Check for cutoff @@ -409,8 +428,9 @@ def bremsstrahlung(reaction, particle_container, collision_data_container, mcdc, @njit def ionization( - reaction, particle_container, collision_data_container, element, mcdc, data + reaction, particle_container, collision_data_container, element, program, data ): + simulation = util.access_simulation(program) particle = particle_container[0] collision_data = collision_data_container[0] @@ -426,8 +446,8 @@ def ionization( xs_sub_ID = int( mcdc_get.electron_ionization_reaction.subshell_x_IDs(i, reaction, data) ) - xs_sub_table = mcdc["data"][xs_sub_ID] - xs_sub_i = evaluate_data(E, xs_sub_table, mcdc, data) + xs_sub_table = simulation["data"][xs_sub_ID] + xs_sub_i = evaluate_data(E, xs_sub_table, simulation, data) xs_vals[i] = xs_sub_i total += xs_sub_i @@ -456,10 +476,8 @@ def ionization( chosen, reaction, data ) ) - dist_base = mcdc["distributions"][dist_ID] - T_delta = sample_distribution( - E, dist_base, particle_container, mcdc, data, scale=False - ) + dist_base = simulation["distributions"][dist_ID] + T_delta = sample_distribution(E, dist_base, particle_container, simulation, data) # Primary outgoing energy E_out = E - B - T_delta @@ -516,7 +534,7 @@ def ionization( particle_new["uz"] = uz_delta particle_new["w"] = particle["w"] - particle_bank_module.bank_active_particle(particle_container_new, mcdc) + particle_bank_module.bank_active_particle(particle_container_new, program) @njit diff --git a/mcdc/transport/physics/interface.py b/mcdc/transport/physics/interface.py index f5772e28..ef3cef34 100644 --- a/mcdc/transport/physics/interface.py +++ b/mcdc/transport/physics/interface.py @@ -16,12 +16,12 @@ @njit -def particle_speed(particle_container, mcdc, data): +def particle_speed(particle_container, simulation, data): particle = particle_container[0] if particle["particle_type"] == PARTICLE_NEUTRON: - return neutron.particle_speed(particle_container, mcdc, data) + return neutron.particle_speed(particle_container, simulation, data) elif particle["particle_type"] == PARTICLE_ELECTRON: - return electron.particle_speed(particle_container, mcdc, data) + return electron.particle_speed(particle_container, simulation, data) return -1.0 @@ -31,21 +31,21 @@ def particle_speed(particle_container, mcdc, data): @njit -def macro_xs(reaction_type, particle_container, mcdc, data): +def macro_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] if particle["particle_type"] == PARTICLE_NEUTRON: - return neutron.macro_xs(reaction_type, particle_container, mcdc, data) + return neutron.macro_xs(reaction_type, particle_container, simulation, data) elif particle["particle_type"] == PARTICLE_ELECTRON: - return electron.macro_xs(reaction_type, particle_container, mcdc, data) + return electron.macro_xs(reaction_type, particle_container, simulation, data) return -1.0 @njit -def neutron_production_xs(reaction_type, particle_container, mcdc, data): +def neutron_production_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] if particle["particle_type"] == PARTICLE_NEUTRON: return neutron.neutron_production_xs( - reaction_type, particle_container, mcdc, data + reaction_type, particle_container, simulation, data ) return -1.0 @@ -56,15 +56,15 @@ def neutron_production_xs(reaction_type, particle_container, mcdc, data): @njit -def collision_distance(particle_container, mcdc, data): +def collision_distance(particle_container, simulation, data): particle = particle_container[0] # Get total cross-section SigmaT = 0.0 if particle["particle_type"] == PARTICLE_NEUTRON: - SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, mcdc, data) + SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, simulation, data) elif particle["particle_type"] == PARTICLE_ELECTRON: - SigmaT = macro_xs(ELECTRON_REACTION_TOTAL, particle_container, mcdc, data) + SigmaT = macro_xs(ELECTRON_REACTION_TOTAL, particle_container, simulation, data) # Vacuum material? if SigmaT == 0.0: @@ -77,10 +77,10 @@ def collision_distance(particle_container, mcdc, data): @njit -def collision(particle_container, collision_data_container, mcdc, data): +def collision(particle_container, collision_data_container, program, data): particle = particle_container[0] if particle["particle_type"] == PARTICLE_NEUTRON: - neutron.collision(particle_container, collision_data_container, mcdc, data) + neutron.collision(particle_container, collision_data_container, program, data) elif particle["particle_type"] == PARTICLE_ELECTRON: - electron.collision(particle_container, collision_data_container, mcdc, data) + electron.collision(particle_container, collision_data_container, program, data) diff --git a/mcdc/transport/physics/neutron/interface.py b/mcdc/transport/physics/neutron/interface.py index 0264745e..29f78602 100644 --- a/mcdc/transport/physics/neutron/interface.py +++ b/mcdc/transport/physics/neutron/interface.py @@ -2,7 +2,9 @@ #### +import mcdc.transport.physics.neutron.multigroup as multigroup import mcdc.transport.physics.neutron.native as native +import mcdc.transport.util as util # ====================================================================================== # Particle attributes @@ -10,8 +12,11 @@ @njit -def particle_speed(particle_container, mcdc, data): - return native.particle_speed(particle_container) +def particle_speed(particle_container, simulation, data): + if simulation["settings"]["neutron_multigroup_mode"]: + return multigroup.particle_speed(particle_container, simulation, data) + else: + return native.particle_speed(particle_container) # ====================================================================================== @@ -20,13 +25,23 @@ def particle_speed(particle_container, mcdc, data): @njit -def macro_xs(reaction_type, particle_container, mcdc, data): - return native.macro_xs(reaction_type, particle_container, mcdc, data) +def macro_xs(reaction_type, particle_container, simulation, data): + if simulation["settings"]["neutron_multigroup_mode"]: + return multigroup.macro_xs(reaction_type, particle_container, simulation, data) + else: + return native.macro_xs(reaction_type, particle_container, simulation, data) @njit -def neutron_production_xs(reaction_type, particle_container, mcdc, data): - return native.neutron_production_xs(reaction_type, particle_container, mcdc, data) +def neutron_production_xs(reaction_type, particle_container, simulation, data): + if simulation["settings"]["neutron_multigroup_mode"]: + return multigroup.neutron_production_xs( + reaction_type, particle_container, simulation, data + ) + else: + return native.neutron_production_xs( + reaction_type, particle_container, simulation, data + ) # ====================================================================================== @@ -35,5 +50,12 @@ def neutron_production_xs(reaction_type, particle_container, mcdc, data): @njit -def collision(particle_container, collision_data_container, mcdc, data): - native.collision(particle_container, collision_data_container, mcdc, data) +def collision(particle_container, collision_data_container, program, data): + simulation = util.access_simulation(program) + + if simulation["settings"]["neutron_multigroup_mode"]: + multigroup.collision( + particle_container, collision_data_container, program, data + ) + else: + native.collision(particle_container, collision_data_container, program, data) diff --git a/mcdc/transport/physics/neutron/multigroup.py b/mcdc/transport/physics/neutron/multigroup.py index 9e24e89a..d5b6c0a5 100644 --- a/mcdc/transport/physics/neutron/multigroup.py +++ b/mcdc/transport/physics/neutron/multigroup.py @@ -10,6 +10,7 @@ import mcdc.transport.particle as particle_module import mcdc.transport.particle_bank as particle_bank_module import mcdc.transport.rng as rng +import mcdc.transport.util as util from mcdc.constant import ( PI, @@ -29,9 +30,9 @@ @njit -def particle_speed(particle_container, mcdc, data): +def particle_speed(particle_container, simulation, data): particle = particle_container[0] - material = mcdc["multigroup_materials"][particle["material_ID"]] + material = simulation["multigroup_materials"][particle["material_ID"]] return mcdc_get.multigroup_material.mgxs_speed(particle["g"], material, data) @@ -41,9 +42,9 @@ def particle_speed(particle_container, mcdc, data): @njit -def macro_xs(reaction_type, particle_container, mcdc, data): +def macro_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] - material = mcdc["multigroup_materials"][particle["material_ID"]] + material = simulation["multigroup_materials"][particle["material_ID"]] g = particle["g"] if reaction_type == NEUTRON_REACTION_TOTAL: @@ -58,42 +59,57 @@ def macro_xs(reaction_type, particle_container, mcdc, data): @njit -def neutron_production_xs(reaction_type, particle_container, mcdc, data): +def neutron_production_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] - material = mcdc["multigroup_materials"][particle["material_ID"]] - + material = simulation["multigroup_materials"][particle["material_ID"]] g = particle["g"] + + # Total production if reaction_type == NEUTRON_REACTION_TOTAL: total = 0.0 - total += neutron_production_xs( - NEUTRON_REACTION_ELASTIC_SCATTERING, - particle_container, - mcdc, - data, - ) - total += neutron_production_xs( - NEUTRON_REACTION_FISSION, particle_container, mcdc, data - ) + + # Scattering production + nu = mcdc_get.multigroup_material.mgxs_nu_s(g, material, data) + xs = mcdc_get.multigroup_material.mgxs_scatter(g, material, data) + total += nu * xs + + # Fission production + nu = mcdc_get.multigroup_material.mgxs_nu_f(g, material, data) + xs = mcdc_get.multigroup_material.mgxs_fission(g, material, data) + total += nu * xs return total + + # Capture production (none) elif reaction_type == NEUTRON_REACTION_CAPTURE: return 0.0 + + # Scattering production elif reaction_type == NEUTRON_REACTION_ELASTIC_SCATTERING: nu = mcdc_get.multigroup_material.mgxs_nu_s(g, material, data) xs = mcdc_get.multigroup_material.mgxs_scatter(g, material, data) return nu * xs + + # Fission production elif reaction_type == NEUTRON_REACTION_FISSION: nu = mcdc_get.multigroup_material.mgxs_nu_f(g, material, data) xs = mcdc_get.multigroup_material.mgxs_fission(g, material, data) return nu * xs + + # Prompt fission production elif reaction_type == NEUTRON_REACTION_FISSION_PROMPT: nu = mcdc_get.multigroup_material.mgxs_nu_p(g, material, data) xs = mcdc_get.multigroup_material.mgxs_fission(g, material, data) return nu * xs + + # Delayed neutron production elif reaction_type == NEUTRON_REACTION_FISSION_DELAYED: nu = mcdc_get.multigroup_material.mgxs_nu_d_total(g, material, data) xs = mcdc_get.multigroup_material.mgxs_fission(g, material, data) return nu * xs + # Unsupported default + return 0.0 + # ====================================================================================== # Collision @@ -101,19 +117,20 @@ def neutron_production_xs(reaction_type, particle_container, mcdc, data): @njit -def collision(particle_container, collision_data_container, mcdc, data): +def collision(particle_container, collision_data_container, program, data): + simulation = util.access_simulation(program) particle = particle_container[0] # Get the reaction cross-sections - SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, mcdc, data) + SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, simulation, data) SigmaS = macro_xs( - NEUTRON_REACTION_ELASTIC_SCATTERING, particle_container, mcdc, data + NEUTRON_REACTION_ELASTIC_SCATTERING, particle_container, simulation, data ) - SigmaC = macro_xs(NEUTRON_REACTION_CAPTURE, particle_container, mcdc, data) - SigmaF = macro_xs(NEUTRON_REACTION_FISSION, particle_container, mcdc, data) + SigmaC = macro_xs(NEUTRON_REACTION_CAPTURE, particle_container, simulation, data) + SigmaF = macro_xs(NEUTRON_REACTION_FISSION, particle_container, simulation, data) # Implicit capture - if mcdc["implicit_capture"]["active"]: + if simulation["implicit_capture"]["active"]: particle["w"] *= (SigmaT - SigmaC) / SigmaT SigmaT -= SigmaC @@ -121,11 +138,11 @@ def collision(particle_container, collision_data_container, mcdc, data): xi = rng.lcg(particle_container) * SigmaT total = SigmaS if total > xi: - scattering(particle_container, mcdc, data) + scattering(particle_container, program, data) else: total += SigmaF if total > xi: - fission(particle_container, mcdc, data) + fission(particle_container, program, data) else: particle["alive"] = False @@ -136,7 +153,9 @@ def collision(particle_container, collision_data_container, mcdc, data): @njit -def scattering(particle_container, mcdc, data): +def scattering(particle_container, program, data): + simulation = util.access_simulation(program) + # Particle attributes particle = particle_container[0] g = particle["g"] @@ -145,7 +164,7 @@ def scattering(particle_container, mcdc, data): uz = particle["uz"] # Material attributes - material = mcdc["multigroup_materials"][particle["material_ID"]] + material = simulation["multigroup_materials"][particle["material_ID"]] G = material["G"] # Kill the current particle @@ -154,8 +173,8 @@ def scattering(particle_container, mcdc, data): # Adjust production and product weights if weighted emission weight_production = 1.0 weight_product = particle["w"] - if mcdc["weighted_emission"]["active"]: - weight_target = mcdc["weighted_emission"]["weight_target"] + if simulation["weighted_emission"]["active"]: + weight_target = simulation["weighted_emission"]["weight_target"] weight_production = particle["w"] / weight_target weight_product = weight_target @@ -164,7 +183,7 @@ def scattering(particle_container, mcdc, data): N = int(math.floor(weight_production * nu_s + rng.lcg(particle_container))) # Set up secondary partice container - particle_container_new = np.zeros(1, type_.particle_data) + particle_container_new = util.local_array(1, type_.particle_data) particle_new = particle_container_new[0] # Create the secondaries @@ -186,7 +205,10 @@ def scattering(particle_container, mcdc, data): particle_new["uz"] = uz_new # Get outgoing spectrum - chi_s = mcdc_get.multigroup_material.mgxs_chi_s_vector(g, material, data) + stride = material["G"] + start = material["mgxs_chi_s_offset"] + g * stride + chi_s = data[start : start + stride] + # Above is equivalent to: chi_s = mcdc_get.multigroup_material.mgxs_chi_s_vector(g, material, data) # Sample outgoing energy xi = rng.lcg(particle_container_new) @@ -207,19 +229,20 @@ def scattering(particle_container, mcdc, data): particle["E"] = particle_new["E"] particle["w"] = particle_new["w"] else: - particle_bank_module.bank_active_particle(particle_container_new, mcdc) + particle_bank_module.bank_active_particle(particle_container_new, program) @njit -def fission(particle_container, mcdc, data): - settings = mcdc["settings"] +def fission(particle_container, program, data): + simulation = util.access_simulation(program) + settings = simulation["settings"] # Particle properties particle = particle_container[0] g = particle["g"] # Material properties - material = mcdc["multigroup_materials"][particle["material_ID"]] + material = simulation["multigroup_materials"][particle["material_ID"]] G = material["G"] J = material["J"] @@ -229,8 +252,8 @@ def fission(particle_container, mcdc, data): # Adjust production and product weights if weighted emission weight_production = 1.0 weight_product = particle["w"] - if mcdc["weighted_emission"]["active"]: - weight_target = mcdc["weighted_emission"]["weight_target"] + if simulation["weighted_emission"]["active"]: + weight_target = simulation["weighted_emission"]["weight_target"] weight_production = particle["w"] / weight_target weight_product = weight_target @@ -238,15 +261,20 @@ def fission(particle_container, mcdc, data): nu = mcdc_get.multigroup_material.mgxs_nu_f(g, material, data) nu_p = mcdc_get.multigroup_material.mgxs_nu_p(g, material, data) if J > 0: - nu_d = mcdc_get.multigroup_material.mgxs_nu_d_vector(g, material, data) + stride = material["J"] + start = material["mgxs_nu_d_offset"] + g * stride + nu_d = data[start : start + stride] + # Above is equivalent to: nu_d = mcdc_get.multigroup_material.mgxs_nu_d_vector(g, material, data) # Get number of secondaries N = int( - math.floor(weight_production * nu / mcdc["k_eff"] + rng.lcg(particle_container)) + math.floor( + weight_production * nu / simulation["k_eff"] + rng.lcg(particle_container) + ) ) # Set up secondary partice container - particle_container_new = np.zeros(1, type_.particle_data) + particle_container_new = util.local_array(1, type_.particle_data) particle_new = particle_container_new[0] # Create the secondaries @@ -268,7 +296,10 @@ def fission(particle_container, mcdc, data): total = nu_p if xi < total: prompt = True - spectrum = mcdc_get.multigroup_material.mgxs_chi_p_vector(g, material, data) + stride = material["G"] + start = material["mgxs_chi_p_offset"] + g * stride + spectrum = data[start : start + stride] + # Above is equivalent to: spectrum = mcdc_get.multigroup_material.mgxs_chi_p_vector(g, material, data) else: prompt = False @@ -276,9 +307,13 @@ def fission(particle_container, mcdc, data): for j in range(J): total += nu_d[j] if xi < total: - spectrum = mcdc_get.multigroup_material.mgxs_chi_d_vector( - j, material, data - ) + stride = material["G"] + start = material["mgxs_chi_d_offset"] + j * stride + spectrum = data[start : start + stride] + # Above is equivalent to: + # spectrum = mcdc_get.multigroup_material.mgxs_chi_d_vector( + # j, material, data + # ) decay = mcdc_get.multigroup_material.mgxs_decay_rate( j, material, data ) @@ -299,8 +334,8 @@ def fission(particle_container, mcdc, data): particle_new["t"] -= math.log(xi) / decay # Eigenvalue mode: bank right away - if settings["eigenvalue_mode"]: - particle_bank_module.bank_census_particle(particle_container_new, mcdc) + if settings["neutron_eigenvalue_mode"]: + particle_bank_module.bank_census_particle(particle_container_new, program) continue # Below is only relevant for fixed-source problem @@ -311,7 +346,7 @@ def fission(particle_container, mcdc, data): # Check if it hits current or next census times hit_current_census = False hit_future_census = False - idx_census = mcdc["idx_census"] + idx_census = simulation["idx_census"] if settings["N_census"] > 1: if particle_new["t"] > mcdc_get.settings.census_time( idx_census, settings, data @@ -335,14 +370,16 @@ def fission(particle_container, mcdc, data): particle["E"] = particle_new["E"] particle["w"] = particle_new["w"] else: - particle_bank_module.bank_active_particle(particle_container_new, mcdc) + particle_bank_module.bank_active_particle( + particle_container_new, program + ) # Hit future census --> add to future bank elif hit_future_census: # Particle will participate in the future - particle_bank_module.bank_future_particle(particle_container_new, mcdc) + particle_bank_module.bank_future_particle(particle_container_new, program) # Hit current census --> add to census bank else: # Particle will participate after the current census is completed - particle_bank_module.bank_census_particle(particle_container_new, mcdc) + particle_bank_module.bank_census_particle(particle_container_new, program) diff --git a/mcdc/transport/physics/neutron/native.py b/mcdc/transport/physics/neutron/native.py index f245ca5b..90f6808f 100644 --- a/mcdc/transport/physics/neutron/native.py +++ b/mcdc/transport/physics/neutron/native.py @@ -1,5 +1,4 @@ import math -import numpy as np from numba import njit @@ -10,6 +9,7 @@ import mcdc.transport.particle as particle_module import mcdc.transport.particle_bank as particle_bank_module import mcdc.transport.rng as rng +import mcdc.transport.util as util from mcdc.constant import ( ANGLE_DISTRIBUTED, @@ -31,8 +31,8 @@ ) from mcdc.transport.data import evaluate_data from mcdc.transport.distribution import ( - sample_correlated_distribution, - sample_distribution, + sample_correlated_distribution_with_scale, + sample_distribution_with_scale, sample_isotropic_cosine, sample_isotropic_direction, sample_multi_table, @@ -70,16 +70,16 @@ def particle_energy_from_speed(speed): @njit -def macro_xs(reaction_type, particle_container, mcdc, data): +def macro_xs(reaction_type, particle_container, simulation, data): particle = particle_container[0] - material = mcdc["native_materials"][particle["material_ID"]] + material = simulation["native_materials"][particle["material_ID"]] E = particle["E"] total = 0.0 for i in range(material["N_nuclide"]): nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) - nuclide = mcdc["nuclides"][nuclide_ID] + nuclide = simulation["nuclides"][nuclide_ID] nuclide_density = mcdc_get.native_material.nuclide_densities(i, material, data) xs = total_micro_xs(reaction_type, E, nuclide, data) @@ -107,6 +107,10 @@ def total_micro_xs(reaction_type, E, nuclide, data): elif reaction_type == NEUTRON_REACTION_FISSION: xs0 = mcdc_get.nuclide.neutron_fission_xs(idx, nuclide, data) xs1 = mcdc_get.nuclide.neutron_fission_xs(idx + 1, nuclide, data) + else: + # Should be unreachable + xs0 = 0.0 + xs1 = 0.0 return linear_interpolation(E, E0, E1, xs0, xs1) @@ -127,89 +131,107 @@ def reaction_micro_xs(E, reaction_base, nuclide, data): @njit -def neutron_production_xs(reaction_type, particle_container, mcdc, data): - particle = particle_container[0] - material_base = mcdc["materials"][particle["material_ID"]] - material = mcdc["native_materials"][material_base["child_ID"]] - +def neutron_production_xs(reaction_type, particle_container, simulation, data): + # Total production if reaction_type == NEUTRON_REACTION_TOTAL: - elastic_type = NEUTRON_REACTION_ELASTIC_SCATTERING - inelastic_type = NEUTRON_REACTION_INELASTIC_SCATTERING - fission_type = NEUTRON_REACTION_FISSION - elastic_xs = neutron_production_xs(elastic_type, particle_container, mcdc, data) - inelastic_xs = neutron_production_xs( - inelastic_type, particle_container, mcdc, data + elastic_xs = macro_xs( + NEUTRON_REACTION_ELASTIC_SCATTERING, particle_container, simulation, data + ) + inelastic_xs = _neutron_inelastic_scattering_production_xs( + particle_container, simulation, data + ) + fission_xs = _neutron_fission_production_xs( + particle_container, simulation, data ) - fission_xs = neutron_production_xs(fission_type, particle_container, mcdc, data) return elastic_xs + inelastic_xs + fission_xs + # Elastic scattering production elif reaction_type == NEUTRON_REACTION_ELASTIC_SCATTERING: - return macro_xs(reaction_type, particle_container, mcdc, data) + return macro_xs(reaction_type, particle_container, simulation, data) + # Capture production (none) elif reaction_type == NEUTRON_REACTION_CAPTURE: return 0.0 + # Inelastic scattering production elif reaction_type == NEUTRON_REACTION_INELASTIC_SCATTERING: - total = 0.0 - for i in range(material["N_nuclide"]): - nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) - nuclide = mcdc["nuclides"][nuclide_ID] + return _neutron_inelastic_scattering_production_xs( + particle_container, simulation, data + ) - E = particle["E"] - nuclide_density = mcdc_get.native_material.nuclide_densities( - i, material, data - ) + # Fission production + elif reaction_type == NEUTRON_REACTION_FISSION: + return _neutron_fission_production_xs(particle_container, simulation, data) - for j in range(nuclide["N_neutron_inelastic_scattering_reaction"]): - reaction_ID = int( - mcdc_get.nuclide.neutron_inelastic_scattering_reaction_IDs( - j, nuclide, data - ) - ) - reaction_base = mcdc["neutron_reactions"][reaction_ID] - reaction = mcdc["neutron_inelastic_scattering_reactions"][ - reaction_base["child_ID"] - ] + # Unsupported default + else: + return 0.0 - xs = reaction_micro_xs(E, reaction_base, nuclide, data) - nu = reaction["multiplicity"] - total += nuclide_density * nu * xs - return total +@njit +def _neutron_inelastic_scattering_production_xs(particle_container, simulation, data): + particle = particle_container[0] + material_base = simulation["materials"][particle["material_ID"]] + material = simulation["native_materials"][material_base["child_ID"]] - elif reaction_type == NEUTRON_REACTION_FISSION: - if not material_base["fissionable"]: - return 0.0 + total = 0.0 + for i in range(material["N_nuclide"]): + nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) + nuclide = simulation["nuclides"][nuclide_ID] - total = 0.0 - for i in range(material["N_nuclide"]): - nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) - nuclide = mcdc["nuclides"][nuclide_ID] - if not nuclide["fissionable"]: - continue + E = particle["E"] + nuclide_density = mcdc_get.native_material.nuclide_densities(i, material, data) - E = particle["E"] - nuclide_density = mcdc_get.native_material.nuclide_densities( - i, material, data + for j in range(nuclide["N_neutron_inelastic_scattering_reaction"]): + reaction_ID = int( + mcdc_get.nuclide.neutron_inelastic_scattering_reaction_IDs( + j, nuclide, data + ) ) + reaction_base = simulation["neutron_reactions"][reaction_ID] + reaction = simulation["neutron_inelastic_scattering_reactions"][ + reaction_base["child_ID"] + ] - for j in range(nuclide["N_neutron_fission_reaction"]): - reaction_ID = int( - mcdc_get.nuclide.neutron_fission_reaction_IDs(j, nuclide, data) - ) - reaction_base = mcdc["neutron_reactions"][reaction_ID] - reaction = mcdc["neutron_fission_reactions"][reaction_base["child_ID"]] + xs = reaction_micro_xs(E, reaction_base, nuclide, data) + nu = reaction["multiplicity"] + total += nuclide_density * nu * xs - xs = reaction_micro_xs(E, reaction_base, nuclide, data) - nu_p = neutron_fission_prompt_multiplicity(E, nuclide, mcdc, data) - nu_d = neutron_fission_delayed_multiplicity(E, nuclide, mcdc, data) - nu = nu_d + nu_p - total += nuclide_density * nu * xs + return total - return total - else: - return -1.0 +@njit +def _neutron_fission_production_xs(particle_container, simulation, data): + particle = particle_container[0] + material_base = simulation["materials"][particle["material_ID"]] + material = simulation["native_materials"][material_base["child_ID"]] + + if not material_base["fissionable"]: + return 0.0 + + total = 0.0 + for i in range(material["N_nuclide"]): + nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) + nuclide = simulation["nuclides"][nuclide_ID] + if not nuclide["fissionable"]: + continue + + E = particle["E"] + nuclide_density = mcdc_get.native_material.nuclide_densities(i, material, data) + + for j in range(nuclide["N_neutron_fission_reaction"]): + reaction_ID = int( + mcdc_get.nuclide.neutron_fission_reaction_IDs(j, nuclide, data) + ) + reaction_base = simulation["neutron_reactions"][reaction_ID] + + xs = reaction_micro_xs(E, reaction_base, nuclide, data) + nu_p = neutron_fission_prompt_multiplicity(E, nuclide, simulation, data) + nu_d = neutron_fission_delayed_multiplicity(E, nuclide, simulation, data) + nu = nu_d + nu_p + total += nuclide_density * nu * xs + + return total # ====================================================================================== @@ -218,10 +240,11 @@ def neutron_production_xs(reaction_type, particle_container, mcdc, data): @njit -def collision(particle_container, collision_data_container, mcdc, data): +def collision(particle_container, collision_data_container, program, data): + simulation = util.access_simulation(program) particle = particle_container[0] collision_data = collision_data_container[0] - material = mcdc["native_materials"][particle["material_ID"]] + material = simulation["native_materials"][particle["material_ID"]] # Particle properties E = particle["E"] @@ -230,12 +253,14 @@ def collision(particle_container, collision_data_container, mcdc, data): # Sample colliding nuclide # ================================================================================== - SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, mcdc, data) + SigmaT = macro_xs(NEUTRON_REACTION_TOTAL, particle_container, simulation, data) # Implicit capture - if mcdc["implicit_capture"]["active"]: + if simulation["implicit_capture"]["active"]: # Calculate capture fraction - SigmaC = macro_xs(NEUTRON_REACTION_CAPTURE, particle_container, mcdc, data) + SigmaC = macro_xs( + NEUTRON_REACTION_CAPTURE, particle_container, simulation, data + ) capture_fraction = SigmaC / SigmaT # Deposit energy captured @@ -244,7 +269,7 @@ def collision(particle_container, collision_data_container, mcdc, data): # Q-value: xs-weighted average over all nuclides and capture reactions for i in range(material["N_nuclide"]): nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) - nuclide = mcdc["nuclides"][nuclide_ID] + nuclide = simulation["nuclides"][nuclide_ID] nuclide_density = mcdc_get.native_material.nuclide_densities( i, material, data ) @@ -252,9 +277,9 @@ def collision(particle_container, collision_data_container, mcdc, data): reaction_ID = int( mcdc_get.nuclide.neutron_capture_reaction_IDs(j, nuclide, data) ) - reaction = mcdc["neutron_capture_reactions"][reaction_ID] + reaction = simulation["neutron_capture_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] xs = reaction_micro_xs(E, reaction_base, nuclide, data) Sigma_rx = nuclide_density * xs collision_data["energy_deposition"] += ( @@ -271,12 +296,12 @@ def collision(particle_container, collision_data_container, mcdc, data): total = 0.0 for i in range(material["N_nuclide"]): nuclide_ID = int(mcdc_get.native_material.nuclide_IDs(i, material, data)) - nuclide = mcdc["nuclides"][nuclide_ID] + nuclide = simulation["nuclides"][nuclide_ID] nuclide_density = mcdc_get.native_material.nuclide_densities(i, material, data) sigmaT = total_micro_xs(NEUTRON_REACTION_TOTAL, E, nuclide, data) - if mcdc["implicit_capture"]["active"]: + if simulation["implicit_capture"]["active"]: sigmaC = total_micro_xs(NEUTRON_REACTION_CAPTURE, E, nuclide, data) sigmaT -= sigmaC @@ -311,9 +336,9 @@ def collision(particle_container, collision_data_container, mcdc, data): i, nuclide, data ) ) - reaction = mcdc["neutron_elastic_scattering_reactions"][reaction_ID] + reaction = simulation["neutron_elastic_scattering_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] total += reaction_micro_xs(E, reaction_base, nuclide, data) # Execute the reaction @@ -323,13 +348,13 @@ def collision(particle_container, collision_data_container, mcdc, data): particle_container, collision_data_container, nuclide, - mcdc, + simulation, data, ) return # Capture - if not mcdc["implicit_capture"]["active"]: + if not simulation["implicit_capture"]["active"]: sigma_capture = total_micro_xs(NEUTRON_REACTION_CAPTURE, E, nuclide, data) total += sigma_capture if xi < total: @@ -339,9 +364,9 @@ def collision(particle_container, collision_data_container, mcdc, data): reaction_ID = int( mcdc_get.nuclide.neutron_capture_reaction_IDs(i, nuclide, data) ) - reaction = mcdc["neutron_capture_reactions"][reaction_ID] + reaction = simulation["neutron_capture_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] xs = reaction_micro_xs(E, reaction_base, nuclide, data) total += xs @@ -352,7 +377,7 @@ def collision(particle_container, collision_data_container, mcdc, data): particle_container, collision_data_container, nuclide, - mcdc, + simulation, data, ) return @@ -368,9 +393,9 @@ def collision(particle_container, collision_data_container, mcdc, data): i, nuclide, data ) ) - reaction = mcdc["neutron_inelastic_scattering_reactions"][reaction_ID] + reaction = simulation["neutron_inelastic_scattering_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] xs = reaction_micro_xs(E, reaction_base, nuclide, data) total += xs @@ -381,7 +406,7 @@ def collision(particle_container, collision_data_container, mcdc, data): particle_container, collision_data_container, nuclide, - mcdc, + program, data, ) return @@ -395,9 +420,9 @@ def collision(particle_container, collision_data_container, mcdc, data): reaction_ID = int( mcdc_get.nuclide.neutron_fission_reaction_IDs(i, nuclide, data) ) - reaction = mcdc["neutron_fission_reactions"][reaction_ID] + reaction = simulation["neutron_fission_reactions"][reaction_ID] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] total += reaction_micro_xs(E, reaction_base, nuclide, data) # Execute the reaction @@ -407,7 +432,7 @@ def collision(particle_container, collision_data_container, mcdc, data): particle_container, collision_data_container, nuclide, - mcdc, + program, data, ) return @@ -420,13 +445,13 @@ def collision(particle_container, collision_data_container, mcdc, data): @njit def capture( - reaction, particle_container, collision_data_container, nuclide, mcdc, data + reaction, particle_container, collision_data_container, nuclide, simulation, data ): particle = particle_container[0] collision_data = collision_data_container[0] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] # Terminate the particle particle["alive"] = False @@ -444,7 +469,7 @@ def capture( @njit def elastic_scattering( - reaction, particle_container, collision_data_container, nuclide, mcdc, data + reaction, particle_container, collision_data_container, nuclide, simulation, data ): particle = particle_container[0] collision_data = collision_data_container[0] @@ -500,7 +525,7 @@ def elastic_scattering( uz = vz / speed # Sample the scattering cosine from the multi-PDF distribution - multi_table = mcdc["multi_table_distributions"][reaction["mu_table_ID"]] + multi_table = simulation["multi_table_distributions"][reaction["mu_table_ID"]] mu0 = sample_multi_table(E, particle_container, multi_table, data) # Scatter the direction in COM @@ -589,13 +614,14 @@ def sample_nucleus_velocity(A, particle_container): @njit def inelastic_scattering( - reaction, particle_container, collision_data_container, nuclide, mcdc, data + reaction, particle_container, collision_data_container, nuclide, program, data ): + simulation = util.access_simulation(program) particle = particle_container[0] collision_data = collision_data_container[0] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] # Particle attributes E = particle["E"] @@ -616,7 +642,7 @@ def inelastic_scattering( use_all_spectrum = N == N_spectrum # Set up secondary partice container - particle_container_new = np.zeros(1, type_.particle_data) + particle_container_new = util.local_array(1, type_.particle_data) particle_new = particle_container_new[0] # Create the secondaries @@ -634,8 +660,8 @@ def inelastic_scattering( elif angle_type == ANGLE_ISOTROPIC: mu = sample_isotropic_cosine(particle_container_new) elif angle_type == ANGLE_DISTRIBUTED: - distribution_base = mcdc["distributions"][reaction["mu_ID"]] - multi_table = mcdc["multi_table_distributions"][ + distribution_base = simulation["distributions"][reaction["mu_ID"]] + multi_table = simulation["multi_table_distributions"][ distribution_base["child_ID"] ] mu = sample_multi_table(E, particle_container_new, multi_table, data) @@ -651,11 +677,15 @@ def inelastic_scattering( n, reaction, data ) ) - spectrum_base = mcdc["distributions"][ID] + spectrum_base = simulation["distributions"][ID] else: - probability_grid = mcdc_get.neutron_inelastic_scattering_reaction.spectrum_probability_grid_all( - reaction, data - ) + offset = reaction["spectrum_probability_grid_offset"] + length = reaction["spectrum_probability_grid_length"] + probability_grid = data[offset : offset + length] + # Above is equivalent to: + # probability_grid = mcdc_get.neutron_inelastic_scattering_reaction.spectrum_probability_grid_all( + # reaction, data + # ) probability_idx = find_bin(E, probability_grid) xi = rng.lcg(particle_container_new) total = 0.0 @@ -672,24 +702,24 @@ def inelastic_scattering( j, reaction, data ) ) - spectrum_base = mcdc["distributions"][ID] + spectrum_base = simulation["distributions"][ID] break # Sample energy if not angle_type == ANGLE_ENERGY_CORRELATED: - E_new = sample_distribution( - E, spectrum_base, particle_container_new, mcdc, data, scale=True + E_new = sample_distribution_with_scale( + E, spectrum_base, particle_container_new, simulation, data ) else: - E_new, mu = sample_correlated_distribution( - E, spectrum_base, particle_container_new, mcdc, data, scale=True + E_new, mu = sample_correlated_distribution_with_scale( + E, spectrum_base, particle_container_new, simulation, data ) # ============================================================================== # Frame transformation # ============================================================================== - reaction_base = mcdc["neutron_reactions"][int(reaction["parent_ID"])] + reaction_base = simulation["neutron_reactions"][int(reaction["parent_ID"])] reference_frame = reaction_base["reference_frame"] if reference_frame == REFERENCE_FRAME_COM: A = nuclide["atomic_weight_ratio"] @@ -725,7 +755,7 @@ def inelastic_scattering( particle["uz"] = particle_new["uz"] particle["E"] = particle_new["E"] else: - particle_bank_module.bank_active_particle(particle_container_new, mcdc) + particle_bank_module.bank_active_particle(particle_container_new, program) # ====================================================================================== @@ -735,15 +765,16 @@ def inelastic_scattering( @njit def fission( - reaction, particle_container, collision_data_container, nuclide, mcdc, data + reaction, particle_container, collision_data_container, nuclide, program, data ): + simulation = util.access_simulation(program) particle = particle_container[0] collision_data = collision_data_container[0] - settings = mcdc["settings"] + settings = simulation["settings"] reaction_base_ID = reaction["parent_ID"] - reaction_base = mcdc["neutron_reactions"][reaction_base_ID] + reaction_base = simulation["neutron_reactions"][reaction_base_ID] # Particle properties E = particle["E"] @@ -762,24 +793,26 @@ def fission( # Adjust production and product weights if weighted emission weight_production = 1.0 weight_product = particle["w"] - if mcdc["weighted_emission"]["active"]: - weight_target = mcdc["weighted_emission"]["weight_target"] + if simulation["weighted_emission"]["active"]: + weight_target = simulation["weighted_emission"]["weight_target"] weight_production = particle["w"] / weight_target weight_product = weight_target # Fission yields N_delayed = nuclide["N_neutron_fission_delayed_precursor"] - nu_p = neutron_fission_prompt_multiplicity(E, nuclide, mcdc, data) - nu_d = neutron_fission_delayed_multiplicity(E, nuclide, mcdc, data) + nu_p = neutron_fission_prompt_multiplicity(E, nuclide, simulation, data) + nu_d = neutron_fission_delayed_multiplicity(E, nuclide, simulation, data) nu = nu_p + nu_d # Get number of secondaries N = int( - math.floor(weight_production * nu / mcdc["k_eff"] + rng.lcg(particle_container)) + math.floor( + weight_production * nu / simulation["k_eff"] + rng.lcg(particle_container) + ) ) # Set up secondary partice container - particle_container_new = np.zeros(1, type_.particle_data) + particle_container_new = util.local_array(1, type_.particle_data) particle_new = particle_container_new[0] # Create the secondaries @@ -819,25 +852,33 @@ def fission( elif angle_type == ANGLE_ISOTROPIC: mu = sample_isotropic_cosine(particle_container_new) elif angle_type == ANGLE_DISTRIBUTED: - distribution_base = mcdc["distributions"][reaction["mu_ID"]] - multi_table = mcdc["multi_table_distributions"][ + distribution_base = simulation["distributions"][reaction["mu_ID"]] + multi_table = simulation["multi_table_distributions"][ distribution_base["child_ID"] ] mu = sample_multi_table(E, particle_container_new, multi_table, data) # Sample energy (also angle if correlated) - spectrum_base = mcdc["distributions"][reaction["spectrum_ID"]] + spectrum_base = simulation["distributions"][reaction["spectrum_ID"]] if not angle_type == ANGLE_ENERGY_CORRELATED: - E_new = sample_distribution( - E, spectrum_base, particle_container_new, mcdc, data, scale=True + E_new = sample_distribution_with_scale( + E, + spectrum_base, + particle_container_new, + simulation, + data, ) else: - E_new, mu = sample_correlated_distribution( - E, spectrum_base, particle_container_new, mcdc, data, scale=True + E_new, mu = sample_correlated_distribution_with_scale( + E, + spectrum_base, + particle_container_new, + simulation, + data, ) # Frame transformation - reaction_base = mcdc["neutron_reactions"][int(reaction["parent_ID"])] + reaction_base = simulation["neutron_reactions"][int(reaction["parent_ID"])] reference_frame = reaction_base["reference_frame"] if reference_frame == REFERENCE_FRAME_COM: A = nuclide["atomic_weight_ratio"] @@ -883,8 +924,8 @@ def fission( # ============================================================================== # Eigenvalue mode: bank right away - if settings["eigenvalue_mode"]: - particle_bank_module.bank_census_particle(particle_container_new, mcdc) + if settings["neutron_eigenvalue_mode"]: + particle_bank_module.bank_census_particle(particle_container_new, program) continue # Below is only relevant for fixed-source problem @@ -895,7 +936,7 @@ def fission( # Check if it hits current or next census times hit_current_census = False hit_future_census = False - idx_census = mcdc["idx_census"] + idx_census = simulation["idx_census"] if settings["N_census"] > 1: if particle_new["t"] > mcdc_get.settings.census_time( idx_census, settings, data @@ -919,26 +960,28 @@ def fission( particle["E"] = particle_new["E"] particle["w"] = particle_new["w"] else: - particle_bank_module.bank_active_particle(particle_container_new, mcdc) + particle_bank_module.bank_active_particle( + particle_container_new, program + ) # Hit future census --> add to future bank elif hit_future_census: # Particle will participate in the future - particle_bank_module.bank_future_particle(particle_container_new, mcdc) + particle_bank_module.bank_future_particle(particle_container_new, program) # Hit current census --> add to census bank else: # Particle will participate after the current census is completed - particle_bank_module.bank_census_particle(particle_container_new, mcdc) + particle_bank_module.bank_census_particle(particle_container_new, program) @njit -def neutron_fission_prompt_multiplicity(E, nuclide, mcdc, data): - data_base = mcdc["data"][nuclide["neutron_fission_prompt_multiplicity_ID"]] - return evaluate_data(E, data_base, mcdc, data) +def neutron_fission_prompt_multiplicity(E, nuclide, simulation, data): + data_base = simulation["data"][nuclide["neutron_fission_prompt_multiplicity_ID"]] + return evaluate_data(E, data_base, simulation, data) @njit -def neutron_fission_delayed_multiplicity(E, nuclide, mcdc, data): - data_base = mcdc["data"][nuclide["neutron_fission_delayed_multiplicity_ID"]] - return evaluate_data(E, data_base, mcdc, data) +def neutron_fission_delayed_multiplicity(E, nuclide, simulation, data): + data_base = simulation["data"][nuclide["neutron_fission_delayed_multiplicity_ID"]] + return evaluate_data(E, data_base, simulation, data) diff --git a/mcdc/transport/physics/util.py b/mcdc/transport/physics/util.py index 1a6fe804..3475a151 100644 --- a/mcdc/transport/physics/util.py +++ b/mcdc/transport/physics/util.py @@ -11,7 +11,11 @@ @njit def evaluate_neutron_xs_energy_grid(e, nuclide, data): - energy_grid = mcdc_get.nuclide.neutron_xs_energy_grid_all(nuclide, data) + offset = nuclide["neutron_xs_energy_grid_offset"] + length = nuclide["neutron_xs_energy_grid_length"] + energy_grid = data[offset : offset + length] + # Above is equivalent to: energy_grid = mcdc_get.nuclide.neutron_xs_energy_grid_all(nuclide, data) + idx = find_bin(e, energy_grid) e0 = energy_grid[idx] e1 = energy_grid[idx + 1] diff --git a/mcdc/transport/simulation.py b/mcdc/transport/simulation.py index d40bc98e..6f2e3f06 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -15,6 +15,7 @@ import mcdc.transport.rng as rng import mcdc.transport.tally as tally_module import mcdc.transport.technique as technique +import mcdc.transport.util as util from mcdc.constant import * from mcdc.print_ import ( @@ -29,13 +30,13 @@ # ====================================================================================== -def fixed_source_simulation(mcdc_arr, data): +def fixed_source_simulation(simulation_container, data): # Ensure `mcdc` exist for the lifetime of the program by intentionally leaking their memory - # adapt.leak(mcdc_arr) - mcdc = mcdc_arr[0] + # adapt.leak(simulation_container) + simulation = simulation_container[0] # Get some settings - settings = mcdc["settings"] + settings = simulation["settings"] N_batch = settings["N_batch"] N_particle = settings["N_particle"] N_census = settings["N_census"] @@ -43,11 +44,11 @@ def fixed_source_simulation(mcdc_arr, data): # Loop over batches for idx_batch in range(N_batch): - mcdc["idx_batch"] = idx_batch + simulation["idx_batch"] = idx_batch seed_batch = rng.split_seed(uint64(idx_batch), settings["rng_seed"]) # Distribute work - mpi.distribute_work(N_particle, mcdc) + mpi.distribute_work(N_particle, simulation) # Print multi-batch header if N_batch > 1: @@ -56,104 +57,104 @@ def fixed_source_simulation(mcdc_arr, data): # Loop over time censuses for idx_census in range(N_census): - mcdc["idx_census"] = idx_census + simulation["idx_census"] = idx_census seed_census = rng.split_seed(uint64(seed_batch), rng.SEED_SPLIT_CENSUS) # Reset tally time filters if census-based tally is used if use_census_based_tally: - tally_module.filter.set_census_based_time_grid(mcdc, data) + tally_module.filter.set_census_based_time_grid(simulation, data) # Accordingly promote future particles to censused particles - if particle_bank_module.get_bank_size(mcdc["bank_future"]) > 0: - particle_bank_module.promote_future_particles(mcdc, data) + if particle_bank_module.get_bank_size(simulation["bank_future"]) > 0: + particle_bank_module.promote_future_particles(simulation, data) # Loop over source particles seed_source = rng.split_seed(uint64(seed_census), rng.SEED_SPLIT_SOURCE) - source_loop(uint64(seed_source), mcdc, data) + source_loop(uint64(seed_source), simulation, data) # Manage particle banks: population control and work rebalance - particle_bank_module.manage_particle_banks(mcdc) + particle_bank_module.manage_particle_banks(simulation) # Time census-based tally closeout if use_census_based_tally: - tally_module.closeout.reduce(mcdc, data) - tally_module.closeout.accumulate(mcdc, data) - if mcdc["mpi_master"]: + tally_module.closeout.reduce(simulation, data) + tally_module.closeout.accumulate(simulation, data) + if simulation["mpi_master"]: with objmode(): - output_module.generate_census_based_tally(mcdc, data) - tally_module.closeout.reset_sum_bins(mcdc, data) + output_module.generate_census_based_tally(simulation, data) + tally_module.closeout.reset_sum_bins(simulation, data) # Terminate census loop if all banks are empty if ( idx_census > 0 - and particle_bank_module.total_size(mcdc["bank_source"]) == 0 - and particle_bank_module.total_size(mcdc["bank_census"]) == 0 - and particle_bank_module.total_size(mcdc["bank_future"]) == 0 + and particle_bank_module.total_size(simulation["bank_source"]) == 0 + and particle_bank_module.total_size(simulation["bank_census"]) == 0 + and particle_bank_module.total_size(simulation["bank_future"]) == 0 ): break # Multi-batch closeout if N_batch > 1: # Reset banks - particle_bank_module.set_bank_size(mcdc["bank_active"], 0) - particle_bank_module.set_bank_size(mcdc["bank_census"], 0) - particle_bank_module.set_bank_size(mcdc["bank_source"], 0) - particle_bank_module.set_bank_size(mcdc["bank_future"], 0) + particle_bank_module.set_bank_size(simulation["bank_active"], 0) + particle_bank_module.set_bank_size(simulation["bank_census"], 0) + particle_bank_module.set_bank_size(simulation["bank_source"], 0) + particle_bank_module.set_bank_size(simulation["bank_future"], 0) if not use_census_based_tally: # Tally history closeout - tally_module.closeout.reduce(mcdc, data) - tally_module.closeout.accumulate(mcdc, data) + tally_module.closeout.reduce(simulation, data) + tally_module.closeout.accumulate(simulation, data) # Tally closeout if not use_census_based_tally: - tally_module.closeout.finalize(mcdc, data) + tally_module.closeout.finalize(simulation, data) -def eigenvalue_simulation(mcdc_arr, data): +def eigenvalue_simulation(simulation_container, data): # Ensure `mcdc` exist for the lifetime of the program # by intentionally leaking their memory - # adapt.leak(mcdc_arr) - mcdc = mcdc_arr[0] + # adapt.leak(simulation_container) + simulation = simulation_container[0] # Get some settings - settings = mcdc["settings"] + settings = simulation["settings"] N_inactive = settings["N_inactive"] N_cycle = settings["N_cycle"] N_particle = settings["N_particle"] # Distribute work - mpi.distribute_work(N_particle, mcdc) + mpi.distribute_work(N_particle, simulation) # Loop over power iteration cycles for idx_cycle in range(N_cycle): - mcdc["idx_cycle"] = idx_cycle + simulation["idx_cycle"] = idx_cycle seed_cycle = rng.split_seed(uint64(idx_cycle), settings["rng_seed"]) # Loop over source particles - source_loop(uint64(seed_cycle), mcdc, data) + source_loop(uint64(seed_cycle), simulation, data) # Tally "history" closeout - tally_module.closeout.eigenvalue_cycle(mcdc, data) - if mcdc["cycle_active"]: - tally_module.closeout.reduce(mcdc, data) - tally_module.closeout.accumulate(mcdc, data) + tally_module.closeout.eigenvalue_cycle(simulation, data) + if simulation["cycle_active"]: + tally_module.closeout.reduce(simulation, data) + tally_module.closeout.accumulate(simulation, data) # Manage particle banks: population control and work rebalance - particle_bank_module.manage_particle_banks(mcdc) + particle_bank_module.manage_particle_banks(simulation) # Print progress with objmode(): - print_progress_eigenvalue(mcdc, data) + print_progress_eigenvalue(simulation, data) # Entering active cycle? - mcdc["idx_cycle"] += 1 - if mcdc["idx_cycle"] >= N_inactive: - mcdc["cycle_active"] = True + simulation["idx_cycle"] += 1 + if simulation["idx_cycle"] >= N_inactive: + simulation["cycle_active"] = True # Tally closeout - tally_module.closeout.finalize(mcdc, data) - tally_module.closeout.eigenvalue_simulation(mcdc) + tally_module.closeout.finalize(simulation, data) + tally_module.closeout.eigenvalue_simulation(simulation) # ============================================================================= @@ -162,41 +163,44 @@ def eigenvalue_simulation(mcdc_arr, data): @njit -def source_loop(seed, mcdc, data): +def source_loop(seed, simulation, data): # Progress bar indicator N_prog = 0 # Loop over particle sources - work_start = mcdc["mpi_work_start"] - work_size = mcdc["mpi_work_size"] + work_start = simulation["mpi_work_start"] + work_size = simulation["mpi_work_size"] for idx_work in range(work_size): - mcdc["idx_work"] = work_start + idx_work - generate_source_particle(work_start, idx_work, seed, mcdc, data) + simulation["idx_work"] = work_start + idx_work + generate_source_particle(work_start, idx_work, seed, simulation, data) # Run the source particle and its secondaries - exhaust_active_bank(mcdc, data) + exhaust_active_bank(simulation, data) - source_closeout(mcdc, idx_work, N_prog, data) + source_closeout(simulation, idx_work, N_prog, data) @njit -def generate_source_particle(work_start, idx_work, seed, mcdc, data): +def generate_source_particle(work_start, idx_work, seed, program, data): """Get a source particle and put into one of the banks""" - settings = mcdc["settings"] - - particle_container = np.zeros(1, type_.particle_data) - particle = particle_container[0] + simulation = util.access_simulation(program) + settings = simulation["settings"] # Get from fixed-source? - if particle_bank_module.get_bank_size(mcdc["bank_source"]) == 0: + if particle_bank_module.get_bank_size(simulation["bank_source"]) == 0: + particle_container = util.local_array(1, type_.particle_data) + particle = particle_container[0] + # Sample source seed_work = rng.split_seed(work_start + idx_work, seed) - source_particle(particle_container, seed_work, mcdc, data) + source_particle(particle_container, seed_work, simulation, data) # Get from source bank else: - particle_container = mcdc["bank_source"]["particles"][idx_work : (idx_work + 1)] + particle_container = simulation["bank_source"]["particle_data"][ + idx_work : (idx_work + 1) + ] particle = particle_container[0] # Skip if beyond time boundary @@ -206,7 +210,7 @@ def generate_source_particle(work_start, idx_work, seed, mcdc, data): # Check if it is beyond current or next census times hit_census = False hit_next_census = False - idx_census = mcdc["idx_census"] + idx_census = simulation["idx_census"] if idx_census < settings["N_census"] - 1: if particle["t"] > mcdc_get.settings.census_time( @@ -219,49 +223,45 @@ def generate_source_particle(work_start, idx_work, seed, mcdc, data): # Put into the right bank if not hit_census: - particle_bank_module.bank_active_particle(particle_container, mcdc) + particle_bank_module.bank_active_particle(particle_container, program) elif not hit_next_census: # Particle will participate after the current census - particle_bank_module.bank_census_particle(particle_container, mcdc) + particle_bank_module.bank_census_particle(particle_container, program) else: # Particle will participate in the future - particle_bank_module.bank_future_particle(particle_container, mcdc) + particle_bank_module.bank_future_particle(particle_container, program) @njit -def exhaust_active_bank(mcdc, data): - particle_container = np.zeros(1, type_.particle) +def exhaust_active_bank(simulation, data): + particle_container = util.local_array(1, type_.particle) particle = particle_container[0] # Loop until active bank is exhausted - while particle_bank_module.get_bank_size(mcdc["bank_active"]) > 0: + while particle_bank_module.get_bank_size(simulation["bank_active"]) > 0: # Get particle from active bank - particle_bank_module.pop_particle(particle_container, mcdc["bank_active"]) - - prep_particle(particle_container, mcdc) + particle_bank_module.pop_particle(particle_container, simulation["bank_active"]) # Particle loop - particle_loop(particle_container, mcdc, data) - - -@njit -def prep_particle(particle_container, mcdc): - particle = particle_container[0] + particle_loop(particle_container, simulation, data) @njit -def source_closeout(mcdc, idx_work, N_prog, data): +def source_closeout(simulation, idx_work, N_prog, data): # Tally history closeout for one-batch fixed-source simulation - if not mcdc["settings"]["eigenvalue_mode"] and mcdc["settings"]["N_batch"] == 1: - if not mcdc["settings"]["use_census_based_tally"]: - tally_module.closeout.accumulate(mcdc, data) + if ( + not simulation["settings"]["neutron_eigenvalue_mode"] + and simulation["settings"]["N_batch"] == 1 + ): + if not simulation["settings"]["use_census_based_tally"]: + tally_module.closeout.accumulate(simulation, data) # Progress printout - percent = (idx_work + 1.0) / mcdc["mpi_work_size"] - if mcdc["settings"]["use_progress_bar"] and int(percent * 100.0) > N_prog: + percent = (idx_work + 1.0) / simulation["mpi_work_size"] + if simulation["settings"]["use_progress_bar"] and int(percent * 100.0) > N_prog: N_prog += 1 with objmode(): - print_progress(percent, mcdc) + print_progress(percent, simulation) # ====================================================================================== @@ -270,19 +270,20 @@ def source_closeout(mcdc, idx_work, N_prog, data): @njit -def particle_loop(particle_container, mcdc, data): +def particle_loop(particle_container, simulation, data): particle = particle_container[0] while particle["alive"]: - step_particle(particle_container, mcdc, data) + step_particle(particle_container, simulation, data) @njit -def step_particle(particle_container, mcdc, data): +def step_particle(particle_container, program, data): + simulation = util.access_simulation(program) particle = particle_container[0] # Determine and move to event - move_to_event(particle_container, mcdc, data) + move_to_event(particle_container, simulation, data) # Execute events if particle["event"] == EVENT_LOST: @@ -293,44 +294,52 @@ def step_particle(particle_container, mcdc, data): collision_data_container = np.zeros(1, type_.collision_data) # Execute the physics - physics.collision(particle_container, collision_data_container, mcdc, data) + physics.collision(particle_container, collision_data_container, program, data) # Score collision tallies - if mcdc["cycle_active"]: + if simulation["cycle_active"]: # Cell tallies - cell = mcdc["cells"][particle["cell_ID"]] + cell = simulation["cells"][particle["cell_ID"]] for i in range(cell["N_tally"]): tally_base_ID = int(mcdc_get.cell.tally_IDs(i, cell, data)) - tally_base = mcdc["tallies"][tally_base_ID] + tally_base = simulation["tallies"][tally_base_ID] # Skip non-collision tallies if tally_base["child_type"] != TALLY_COLLISION: continue - tally = mcdc["collision_tallies"][tally_base["child_ID"]] + tally = simulation["collision_tallies"][tally_base["child_ID"]] tally_module.score.collision_tally( - particle_container, collision_data_container, tally, mcdc, data + particle_container, + collision_data_container, + tally, + simulation, + data, ) # Other collision tallies - for i in range(mcdc["N_collision_tally"]): - tally = mcdc["collision_tallies"][i] + for i in range(simulation["N_collision_tally"]): + tally = simulation["collision_tallies"][i] # Skip cell tallies if tally["spatial_filter_type"] == SPATIAL_FILTER_CELL: continue tally_module.score.collision_tally( - particle_container, collision_data_container, tally, mcdc, data + particle_container, + collision_data_container, + tally, + simulation, + data, ) # Surface and domain crossing if particle["event"] & EVENT_SURFACE_CROSSING: - geometry.surface_crossing(particle_container, mcdc, data) + geometry.surface_crossing(particle_container, simulation, data) # Census time crossing if particle["event"] & EVENT_TIME_CENSUS: - particle_bank_module.bank_census_particle(particle_container, mcdc) + particle_bank_module.bank_census_particle(particle_container, program) particle["alive"] = False # Time boundary crossing @@ -339,12 +348,12 @@ def step_particle(particle_container, mcdc, data): # Weight roulette if particle["alive"]: - technique.weight_roulette(particle_container, mcdc) + technique.weight_roulette(particle_container, simulation) @njit -def move_to_event(particle_container, mcdc, data): - settings = mcdc["settings"] +def move_to_event(particle_container, simulation, data): + settings = simulation["settings"] # ================================================================================== # Preparation (as needed) @@ -353,10 +362,10 @@ def move_to_event(particle_container, mcdc, data): # Multigroup preparation # In MG mode, particle speed is material-dependent. - if settings["multigroup_mode"]: + if settings["neutron_multigroup_mode"]: # If material is not identified yet, locate the particle if particle["material_ID"] == -1: - if not geometry.locate_particle(particle_container, mcdc, data): + if not geometry.locate_particle(particle_container, simulation, data): # Particle is lost particle["event"] = EVENT_LOST return @@ -369,7 +378,7 @@ def move_to_event(particle_container, mcdc, data): # - Set particle boundary event (surface or lattice crossing, or lost) # - Return distance to boundary (surface or lattice) - d_boundary = geometry.inspect_geometry(particle_container, mcdc, data) + d_boundary = geometry.inspect_geometry(particle_container, simulation, data) # Particle is lost? if particle["event"] == EVENT_LOST: @@ -380,19 +389,19 @@ def move_to_event(particle_container, mcdc, data): # ================================================================================== # Distance to domain - speed = physics.particle_speed(particle_container, mcdc, data) + speed = physics.particle_speed(particle_container, simulation, data) # Distance to time boundary d_time_boundary = speed * (settings["time_boundary"] - particle["t"]) # Distance to census time - idx = mcdc["idx_census"] + idx = simulation["idx_census"] d_time_census = speed * ( mcdc_get.settings.census_time(idx, settings, data) - particle["t"] ) # Distance to next collision - d_collision = physics.collision_distance(particle_container, mcdc, data) + d_collision = physics.collision_distance(particle_container, simulation, data) # ================================================================================== # Determine event(s) @@ -428,36 +437,38 @@ def move_to_event(particle_container, mcdc, data): # ================================================================================== # Score tracklength tallies - if mcdc["cycle_active"]: + if simulation["cycle_active"]: # Cell tallies - cell = mcdc["cells"][particle["cell_ID"]] + cell = simulation["cells"][particle["cell_ID"]] for i in range(cell["N_tally"]): tally_base_ID = int(mcdc_get.cell.tally_IDs(i, cell, data)) - tally_base = mcdc["tallies"][tally_base_ID] + tally_base = simulation["tallies"][tally_base_ID] # Skip non-tracklength tallies if tally_base["child_type"] != TALLY_TRACKLENGTH: continue - tally = mcdc["tracklength_tallies"][tally_base["child_ID"]] + tally = simulation["tracklength_tallies"][tally_base["child_ID"]] tally_module.score.tracklength_tally( - particle_container, distance, tally, mcdc, data + particle_container, distance, tally, simulation, data ) # Other tracklength tallies - for i in range(mcdc["N_tracklength_tally"]): - tally = mcdc["tracklength_tallies"][i] + for i in range(simulation["N_tracklength_tally"]): + tally = simulation["tracklength_tallies"][i] # Skip cell tallies if tally["spatial_filter_type"] == SPATIAL_FILTER_CELL: continue tally_module.score.tracklength_tally( - particle_container, distance, tally, mcdc, data + particle_container, distance, tally, simulation, data ) - if settings["eigenvalue_mode"]: - tally_module.score.eigenvalue_tally(particle_container, distance, mcdc, data) + if settings["neutron_eigenvalue_mode"]: + tally_module.score.eigenvalue_tally( + particle_container, distance, simulation, data + ) # Move particle - particle_module.move(particle_container, distance, mcdc, data) + particle_module.move(particle_container, distance, simulation, data) diff --git a/mcdc/transport/source.py b/mcdc/transport/source.py index 776fc9f0..36291b67 100644 --- a/mcdc/transport/source.py +++ b/mcdc/transport/source.py @@ -14,19 +14,19 @@ sample_isotropic_direction, sample_direction, ) -from mcdc.transport.util import find_bin +from mcdc.transport.util import find_bin_with_rules @njit -def source_particle(P_rec_arr, seed, mcdc, data): - P_rec = P_rec_arr[0] - P_rec["rng_seed"] = seed +def source_particle(particle_container, seed, simulation, data): + particle = particle_container[0] + particle["rng_seed"] = seed # Sample source # TODO: use cdf and binary search instead - xi = rng.lcg(P_rec_arr) + xi = rng.lcg(particle_container) tot = 0.0 - for source in mcdc["sources"]: + for source in simulation["sources"]: tot += source["probability"] if tot >= xi: break @@ -37,65 +37,86 @@ def source_particle(P_rec_arr, seed, mcdc, data): y = source["point"][1] z = source["point"][2] else: - x = sample_uniform(source["x"][0], source["x"][1], P_rec_arr) - y = sample_uniform(source["y"][0], source["y"][1], P_rec_arr) - z = sample_uniform(source["z"][0], source["z"][1], P_rec_arr) + x = sample_uniform(source["x"][0], source["x"][1], particle_container) + y = sample_uniform(source["y"][0], source["y"][1], particle_container) + z = sample_uniform(source["z"][0], source["z"][1], particle_container) # Direction if source["isotropic_direction"]: - ux, uy, uz = sample_isotropic_direction(P_rec_arr) + ux, uy, uz = sample_isotropic_direction(particle_container) elif source["white_direction"]: rx = source["direction"][0] ry = source["direction"][1] rz = source["direction"][2] - ux, uy, uz = sample_white_direction(rx, ry, rz, P_rec_arr) + ux, uy, uz = sample_white_direction(rx, ry, rz, particle_container) elif source["mono_direction"]: ux = source["direction"][0] uy = source["direction"][1] uz = source["direction"][2] else: ux, uy, uz = sample_direction( - source["polar_cosine"], source["azimuthal"], source["direction"], P_rec_arr + source["polar_cosine"], + source["azimuthal"], + source["direction"], + particle_container, ) # Energy - if mcdc["settings"]["multigroup_mode"]: + if simulation["settings"]["neutron_multigroup_mode"]: E = 0.0 if source["mono_energetic"]: g = source["energy_group"] else: ID = source["energy_group_pmf_ID"] - pmf = mcdc["pmf_distributions"][ID] - g = sample_pmf(pmf, P_rec_arr, data) + pmf = simulation["pmf_distributions"][ID] + g = sample_pmf(pmf, particle_container, data) else: g = 0 if source["mono_energetic"]: E = source["energy"] else: ID = source["energy_pdf_ID"] - table = mcdc["tabulated_distributions"][ID] - E = sample_tabulated(table, P_rec_arr, data) + table = simulation["tabulated_distributions"][ID] + E = sample_tabulated(table, particle_container, data) # Time if source["discrete_time"]: t = source["time"] else: - t = sample_uniform(source["time_range"][0], source["time_range"][1], P_rec_arr) + t = sample_uniform( + source["time_range"][0], source["time_range"][1], particle_container + ) # Motion translation if source["moving"]: # Get moving interval index wrt the given time - time_grid = mcdc_get.source.move_time_grid_all(source, data) - idx = find_bin(t, time_grid, epsilon=COINCIDENCE_TOLERANCE_TIME, go_lower=False) + time_grid = data[ + source["move_time_grid_offset"] : ( + source["move_time_grid_offset"] + source["N_move_grid"] + ) + ] + # Above is equivalent to: time_grid = mcdc_get.source.move_time_grid_all(source, data) + + tolerance = COINCIDENCE_TOLERANCE_TIME + go_lower = False + idx = find_bin_with_rules(t, time_grid, tolerance, go_lower) # Coinciding cases if abs(time_grid[idx + 1] - t) < COINCIDENCE_TOLERANCE: idx += 1 - # Surface move translations, velocities, and time grid - trans_0 = mcdc_get.surface.move_translations_vector(idx, source, data) - time_0 = mcdc_get.surface.move_time_grid(idx, source, data) - V = mcdc_get.surface.move_velocities_vector(idx, source, data) + # Source move translations + start = source["move_translations_offset"] + idx * 3 + trans_0 = data[start : start + 3] + # Above is equivalent to: trans_0 = mcdc_get.source.move_translations_vector(idx, source, data) + + # Source move velocities + start = source["move_velocities_offset"] + idx * 3 + V = data[start : start + 3] + # Above is equivalent to: V = mcdc_get.source.move_velocities_vector(idx, source, data) + + # Source move time grid + time_0 = mcdc_get.source.move_time_grid(idx, source, data) # Translate the particle t_local = t - time_0 @@ -104,14 +125,14 @@ def source_particle(P_rec_arr, seed, mcdc, data): z += trans_0[2] + V[2] * t_local # Make and return particle - P_rec["x"] = x - P_rec["y"] = y - P_rec["z"] = z - P_rec["t"] = t - P_rec["ux"] = ux - P_rec["uy"] = uy - P_rec["uz"] = uz - P_rec["g"] = g - P_rec["E"] = E - P_rec["w"] = 1.0 - P_rec["particle_type"] = source["particle_type"] + particle["x"] = x + particle["y"] = y + particle["z"] = z + particle["t"] = t + particle["ux"] = ux + particle["uy"] = uy + particle["uz"] = uz + particle["g"] = g + particle["E"] = E + particle["w"] = 1.0 + particle["particle_type"] = source["particle_type"] diff --git a/mcdc/transport/tally/closeout.py b/mcdc/transport/tally/closeout.py index 70743b7b..96f2cd66 100644 --- a/mcdc/transport/tally/closeout.py +++ b/mcdc/transport/tally/closeout.py @@ -26,19 +26,19 @@ @njit -def reduce(mcdc, data): - for tally in mcdc["tallies"]: - _reduce(tally, mcdc, data) +def reduce(simulation, data): + for tally in simulation["tallies"]: + _reduce(tally, simulation, data) @njit -def _reduce(tally, mcdc, data): +def _reduce(tally, simulation, data): N = tally["bin_length"] start = tally["bin_offset"] end = start + N # Normalize - N_particle = mcdc["settings"]["N_particle"] + N_particle = simulation["settings"]["N_particle"] for i in range(N): data[start + i] /= N_particle @@ -55,8 +55,8 @@ def _reduce(tally, mcdc, data): @njit -def accumulate(mcdc, data): - for tally in mcdc["tallies"]: +def accumulate(simulation, data): + for tally in simulation["tallies"]: _accumulate(tally, data) @@ -91,15 +91,15 @@ def _accumulate(tally, data): @njit -def finalize(mcdc, data): - for tally in mcdc["tallies"]: - _finalize(tally, mcdc, data) +def finalize(simulation, data): + for tally in simulation["tallies"]: + _finalize(tally, simulation, data) @njit -def _finalize(tally, mcdc, data): - N_history = mcdc["settings"]["N_particle"] - N_batch = mcdc["settings"]["N_batch"] +def _finalize(tally, simulation, data): + N_history = simulation["settings"]["N_particle"] + N_batch = simulation["settings"]["N_batch"] N_bin = tally["bin_length"] sum_start = tally["bin_sum_offset"] sum_sq_start = tally["bin_sum_square_offset"] @@ -109,8 +109,8 @@ def _finalize(tally, mcdc, data): if N_batch > 1: N_history = N_batch - elif mcdc["settings"]["eigenvalue_mode"]: - N_history = mcdc["settings"]["N_active"] + elif simulation["settings"]["neutron_eigenvalue_mode"]: + N_history = simulation["settings"]["N_active"] else: # MPI Reduce @@ -147,8 +147,8 @@ def _finalize(tally, mcdc, data): @njit -def reset_sum_bins(mcdc, data): - for tally in mcdc["tallies"]: +def reset_sum_bins(simulation, data): + for tally in simulation["tallies"]: _reset_sum_bins(tally, data) @@ -169,9 +169,9 @@ def _reset_sum_bins(tally, data): @njit -def eigenvalue_cycle(mcdc, data): - idx_cycle = mcdc["idx_cycle"] - N_particle = mcdc["settings"]["N_particle"] +def eigenvalue_cycle(simulation, data): + idx_cycle = simulation["idx_cycle"] + N_particle = simulation["settings"]["N_particle"] # MPI Allreduce buff_nuSigmaF = np.zeros(1, np.float64) @@ -181,64 +181,68 @@ def eigenvalue_cycle(mcdc, data): buff_Cmax = np.zeros(1, np.float64) with objmode(): MPI.COMM_WORLD.Allreduce( - np.array(mcdc["eigenvalue_tally_nuSigmaF"]), buff_nuSigmaF, MPI.SUM + np.array(simulation["eigenvalue_tally_nuSigmaF"]), buff_nuSigmaF, MPI.SUM ) - if mcdc["cycle_active"]: + if simulation["cycle_active"]: MPI.COMM_WORLD.Allreduce( - np.array(mcdc["eigenvalue_tally_n"]), buff_n, MPI.SUM + np.array(simulation["eigenvalue_tally_n"]), buff_n, MPI.SUM ) - MPI.COMM_WORLD.Allreduce(np.array([mcdc["n_max"]]), buff_nmax, MPI.MAX) MPI.COMM_WORLD.Allreduce( - np.array(mcdc["eigenvalue_tally_C"]), buff_C, MPI.SUM + np.array([simulation["n_max"]]), buff_nmax, MPI.MAX + ) + MPI.COMM_WORLD.Allreduce( + np.array(simulation["eigenvalue_tally_C"]), buff_C, MPI.SUM + ) + MPI.COMM_WORLD.Allreduce( + np.array([simulation["C_max"]]), buff_Cmax, MPI.MAX ) - MPI.COMM_WORLD.Allreduce(np.array([mcdc["C_max"]]), buff_Cmax, MPI.MAX) # Update and store k_eff - mcdc["k_eff"] = buff_nuSigmaF[0] / N_particle - mcdc_set.simulation.k_cycle(idx_cycle, mcdc, data, value=mcdc["k_eff"]) + simulation["k_eff"] = buff_nuSigmaF[0] / N_particle + mcdc_set.simulation.k_cycle(idx_cycle, simulation, data, value=simulation["k_eff"]) # Normalize other eigenvalue/global tallies tally_n = buff_n[0] / N_particle tally_C = buff_C[0] / N_particle # Maximum densities - mcdc["n_max"] = buff_nmax[0] - mcdc["C_max"] = buff_Cmax[0] + simulation["n_max"] = buff_nmax[0] + simulation["C_max"] = buff_Cmax[0] # Accumulate running average - if mcdc["cycle_active"]: - mcdc["k_avg"] += mcdc["k_eff"] - mcdc["k_sdv"] += mcdc["k_eff"] * mcdc["k_eff"] - mcdc["n_avg"] += tally_n - mcdc["n_sdv"] += tally_n * tally_n - mcdc["C_avg"] += tally_C - mcdc["C_sdv"] += tally_C * tally_C - - N = 1 + mcdc["idx_cycle"] - mcdc["settings"]["N_inactive"] - mcdc["k_avg_running"] = mcdc["k_avg"] / N + if simulation["cycle_active"]: + simulation["k_avg"] += simulation["k_eff"] + simulation["k_sdv"] += simulation["k_eff"] * simulation["k_eff"] + simulation["n_avg"] += tally_n + simulation["n_sdv"] += tally_n * tally_n + simulation["C_avg"] += tally_C + simulation["C_sdv"] += tally_C * tally_C + + N = 1 + simulation["idx_cycle"] - simulation["settings"]["N_inactive"] + simulation["k_avg_running"] = simulation["k_avg"] / N if N == 1: - mcdc["k_sdv_running"] = 0.0 + simulation["k_sdv_running"] = 0.0 else: - mcdc["k_sdv_running"] = math.sqrt( - (mcdc["k_sdv"] / N - mcdc["k_avg_running"] ** 2) / (N - 1) + simulation["k_sdv_running"] = math.sqrt( + (simulation["k_sdv"] / N - simulation["k_avg_running"] ** 2) / (N - 1) ) # Reset accumulators - mcdc["eigenvalue_tally_nuSigmaF"][0] = 0.0 - mcdc["eigenvalue_tally_n"][0] = 0.0 - mcdc["eigenvalue_tally_C"][0] = 0.0 + simulation["eigenvalue_tally_nuSigmaF"][0] = 0.0 + simulation["eigenvalue_tally_n"][0] = 0.0 + simulation["eigenvalue_tally_C"][0] = 0.0 # ===================================================================== # Gyration radius # ===================================================================== - if mcdc["settings"]["use_gyration_radius"]: + if simulation["settings"]["use_gyration_radius"]: # Center of mass - N_local = particle_bank_module.get_bank_size(mcdc["bank_census"]) + N_local = particle_bank_module.get_bank_size(simulation["bank_census"]) total_local = np.zeros(4, np.float64) # [x,y,z,W] total = np.zeros(4, np.float64) for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] total_local[0] += P["x"] * P["w"] total_local[1] += P["y"] * P["w"] total_local[2] += P["z"] * P["w"] @@ -255,10 +259,10 @@ def eigenvalue_cycle(mcdc, data): # Distance RMS rms_local = np.zeros(1, np.float64) rms = np.zeros(1, np.float64) - gr_type = mcdc["settings"]["gyration_radius_type"] + gr_type = simulation["settings"]["gyration_radius_type"] if gr_type == GYRATION_RADIUS_ALL: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ( (P["x"] - com_x) ** 2 + (P["y"] - com_y) ** 2 @@ -266,27 +270,27 @@ def eigenvalue_cycle(mcdc, data): ) * P["w"] elif gr_type == GYRATION_RADIUS_INFINITE_X: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ((P["y"] - com_y) ** 2 + (P["z"] - com_z) ** 2) * P["w"] elif gr_type == GYRATION_RADIUS_INFINITE_Y: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ((P["x"] - com_x) ** 2 + (P["z"] - com_z) ** 2) * P["w"] elif gr_type == GYRATION_RADIUS_INFINITE_Z: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ((P["x"] - com_x) ** 2 + (P["y"] - com_y) ** 2) * P["w"] elif gr_type == GYRATION_RADIUS_ONLY_X: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ((P["x"] - com_x) ** 2) * P["w"] elif gr_type == GYRATION_RADIUS_ONLY_Y: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ((P["y"] - com_y) ** 2) * P["w"] elif gr_type == GYRATION_RADIUS_ONLY_Z: for i in range(N_local): - P = mcdc["bank_census"]["particles"][i] + P = simulation["bank_census"]["particle_data"][i] rms_local[0] += ((P["z"] - com_z) ** 2) * P["w"] # MPI Allreduce @@ -295,17 +299,21 @@ def eigenvalue_cycle(mcdc, data): rms = math.sqrt(rms[0] / W) # Gyration radius - mcdc_set.simulation.gyration_radius(idx_cycle, mcdc, data, value=rms) + mcdc_set.simulation.gyration_radius(idx_cycle, simulation, data, value=rms) @njit -def eigenvalue_simulation(mcdc): - N = mcdc["settings"]["N_active"] - mcdc["n_avg"] /= N - mcdc["C_avg"] /= N +def eigenvalue_simulation(simulation): + N = simulation["settings"]["N_active"] + simulation["n_avg"] /= N + simulation["C_avg"] /= N if N > 1: - mcdc["n_sdv"] = math.sqrt((mcdc["n_sdv"] / N - mcdc["n_avg"] ** 2) / (N - 1)) - mcdc["C_sdv"] = math.sqrt((mcdc["C_sdv"] / N - mcdc["C_avg"] ** 2) / (N - 1)) + simulation["n_sdv"] = math.sqrt( + (simulation["n_sdv"] / N - simulation["n_avg"] ** 2) / (N - 1) + ) + simulation["C_sdv"] = math.sqrt( + (simulation["C_sdv"] / N - simulation["C_avg"] ** 2) / (N - 1) + ) else: - mcdc["n_sdv"] = 0.0 - mcdc["C_sdv"] = 0.0 + simulation["n_sdv"] = 0.0 + simulation["C_sdv"] = 0.0 diff --git a/mcdc/transport/tally/filter.py b/mcdc/transport/tally/filter.py index 22856ba1..a743f785 100644 --- a/mcdc/transport/tally/filter.py +++ b/mcdc/transport/tally/filter.py @@ -12,7 +12,7 @@ COINCIDENCE_TOLERANCE_ENERGY, COINCIDENCE_TOLERANCE_TIME, ) -from mcdc.transport.util import find_bin +from mcdc.transport.util import find_bin_with_tolerance, find_bin_with_rules @njit @@ -55,22 +55,33 @@ def get_direction_index(particle_container, tally, data): azi *= -1 tolerance = COINCIDENCE_TOLERANCE_DIRECTION - i_mu = find_bin(mu, mcdc_get.tally.mu_all(tally, data), tolerance) - i_azi = find_bin(azi, mcdc_get.tally.azi_all(tally, data), tolerance) + + grid_mu = data[tally["mu_offset"] : (tally["mu_offset"] + tally["mu_length"])] + # Above is equivalent to: grid_mu = mcdc_get.tally.mu_all(tally, data) + grid_azi = data[tally["azi_offset"] : (tally["azi_offset"] + tally["azi_length"])] + # Above is equivalent to: grid_azi = mcdc_get.tally.azi_all(tally, data) + + i_mu = find_bin_with_tolerance(mu, grid_mu, tolerance) + i_azi = find_bin_with_tolerance(azi, grid_azi, tolerance) return i_mu, i_azi @njit -def get_energy_index(particle_container, tally, data, multigroup_mode): +def get_energy_index(particle_container, tally, data, neutron_multigroup_mode): particle = particle_container[0] - if multigroup_mode: + if neutron_multigroup_mode: E = particle["g"] else: E = particle["E"] tolerance = COINCIDENCE_TOLERANCE_ENERGY - return find_bin(E, mcdc_get.tally.energy_all(tally, data), tolerance) + grid_energy = data[ + tally["energy_offset"] : (tally["energy_offset"] + tally["energy_length"]) + ] + # Above is equivalent to: grid_energy = mcdc_get.tally.energy_all(tally, data) + + return find_bin_with_tolerance(E, grid_energy, tolerance) @njit @@ -80,17 +91,21 @@ def get_time_index(particle_container, tally, data): # Particle properties time = particle["t"] + grid_time = data[ + tally["time_offset"] : (tally["time_offset"] + tally["time_length"]) + ] + # Above is equivalent to: grid_time = mcdc_get.tally.time_all(tally, data) + tolerance = COINCIDENCE_TOLERANCE_TIME - return find_bin( - time, mcdc_get.tally.time_all(tally, data), tolerance, go_lower=False - ) + go_lower = False + return find_bin_with_rules(time, grid_time, tolerance, go_lower) @njit -def set_census_based_time_grid(mcdc, data): - settings = mcdc["settings"] +def set_census_based_time_grid(simulation, data): + settings = simulation["settings"] tally_frequency = settings["census_tally_frequency"] - idx_census = mcdc["idx_census"] + idx_census = simulation["idx_census"] # Starting time if idx_census == 0: @@ -105,7 +120,7 @@ def set_census_based_time_grid(mcdc, data): dt = (t_end - t_start) / tally_frequency # Set the time grid to all tallies - for tally in mcdc["tallies"]: + for tally in simulation["tallies"]: mcdc_set.tally.time(0, tally, data, t_start) for j in range(tally_frequency): t_next = mcdc_get.tally.time(j, tally, data) + dt diff --git a/mcdc/transport/tally/score.py b/mcdc/transport/tally/score.py index 5e5f0296..924c0338 100644 --- a/mcdc/transport/tally/score.py +++ b/mcdc/transport/tally/score.py @@ -5,6 +5,7 @@ import mcdc.mcdc_get as mcdc_get import mcdc.transport.mesh as mesh_module import mcdc.transport.physics as physics +import mcdc.transport.util as util from mcdc.constant import ( AXIS_T, @@ -28,7 +29,6 @@ ) from mcdc.transport.geometry.surface import get_normal_component from mcdc.transport.tally.filter import get_filter_indices -from mcdc.transport.util import atomic_add # ====================================================================================== # Surface tally @@ -36,12 +36,12 @@ @njit -def surface_tally(particle_container, surface, tally, mcdc, data): +def surface_tally(particle_container, surface, tally, simulation, data): particle = particle_container[0] - tally_base = mcdc["tallies"][tally["parent_ID"]] + tally_base = simulation["tallies"][tally["parent_ID"]] # Get filter indices - MG_mode = mcdc["settings"]["multigroup_mode"] + MG_mode = simulation["settings"]["neutron_multigroup_mode"] i_mu, i_azi, i_energy, i_time = get_filter_indices( particle_container, tally_base, data, MG_mode ) @@ -60,7 +60,7 @@ def surface_tally(particle_container, surface, tally, mcdc, data): ) # Flux - speed = physics.particle_speed(particle_container, mcdc, data) + speed = physics.particle_speed(particle_container, simulation, data) mu = get_normal_component(particle_container, speed, surface, data) flux = particle["w"] / abs(mu) @@ -69,10 +69,10 @@ def surface_tally(particle_container, surface, tally, mcdc, data): score_type = mcdc_get.tally.scores(i_score, tally_base, data) score = 0.0 if score_type == SCORE_NET_CURRENT: - surface = mcdc["surfaces"][particle["surface_ID"]] + surface = simulation["surfaces"][particle["surface_ID"]] mu = get_normal_component(particle_container, speed, surface, data) score = flux * mu - atomic_add(data, idx_base + i_score, score) + util.atomic_add(data, idx_base + i_score, score) # ====================================================================================== @@ -81,13 +81,15 @@ def surface_tally(particle_container, surface, tally, mcdc, data): @njit -def collision_tally(particle_container, collision_data_container, tally, mcdc, data): +def collision_tally( + particle_container, collision_data_container, tally, simulation, data +): particle = particle_container[0] collision_data = collision_data_container[0] - tally_base = mcdc["tallies"][tally["parent_ID"]] + tally_base = simulation["tallies"][tally["parent_ID"]] # Get filter indices - MG_mode = mcdc["settings"]["multigroup_mode"] + MG_mode = simulation["settings"]["neutron_multigroup_mode"] i_mu, i_azi, i_energy, i_time = get_filter_indices( particle_container, tally_base, data, MG_mode ) @@ -100,8 +102,10 @@ def collision_tally(particle_container, collision_data_container, tally, mcdc, d i_x, i_y, i_z = 0, 0, 0 mesh_tally = tally["spatial_filter_type"] == SPATIAL_FILTER_MESH if mesh_tally: - mesh = mcdc["meshes"][tally["spatial_filter_ID"]] - i_x, i_y, i_z = mesh_module.get_indices(particle_container, mesh, mcdc, data) + mesh = simulation["meshes"][tally["spatial_filter_ID"]] + i_x, i_y, i_z = mesh_module.get_indices( + particle_container, mesh, simulation, data + ) # No score outside mesh bins if i_x == -1 or i_y == -1 or i_z == -1: @@ -128,7 +132,7 @@ def collision_tally(particle_container, collision_data_container, tally, mcdc, d score = 0.0 if score_type == SCORE_ENERGY_DEPOSITION: score = collision_data["energy_deposition"] - atomic_add(data, idx_base + i_score, score) + util.atomic_add(data, idx_base + i_score, score) # ====================================================================================== @@ -137,12 +141,12 @@ def collision_tally(particle_container, collision_data_container, tally, mcdc, d @njit -def tracklength_tally(particle_container, distance, tally, mcdc, data): +def tracklength_tally(particle_container, distance, tally, simulation, data): particle = particle_container[0] - tally_base = mcdc["tallies"][tally["parent_ID"]] + tally_base = simulation["tallies"][tally["parent_ID"]] # Get filter indices - MG_mode = mcdc["settings"]["multigroup_mode"] + MG_mode = simulation["settings"]["neutron_multigroup_mode"] i_mu, i_azi, i_energy, i_time = get_filter_indices( particle_container, tally_base, data, MG_mode ) @@ -159,7 +163,7 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): ux = particle["ux"] uy = particle["uy"] uz = particle["uz"] - ut = 1.0 / physics.particle_speed(particle_container, mcdc, data) + ut = 1.0 / physics.particle_speed(particle_container, simulation, data) x_final = x + ux * distance y_final = y + uy * distance z_final = z + uz * distance @@ -189,15 +193,17 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): # Mesh axis indices i_x, i_y, i_z = 0, 0, 0 if mesh_tally: - mesh = mcdc["meshes"][tally["spatial_filter_ID"]] + mesh = simulation["meshes"][tally["spatial_filter_ID"]] # Mesh axis indices - i_x, i_y, i_z = mesh_module.get_indices(particle_container, mesh, mcdc, data) + i_x, i_y, i_z = mesh_module.get_indices( + particle_container, mesh, simulation, data + ) # No score if particle does not cross the mesh bins # Also get the appropriate index if needed - x_min = mesh_module.get_x(0, mesh, mcdc, data) - x_max = mesh_module.get_x(mesh["Nx"], mesh, mcdc, data) + x_min = mesh_module.get_x(0, mesh, simulation, data) + x_max = mesh_module.get_x(mesh["Nx"], mesh, simulation, data) if ux > 0.0: if ( x_final < x_min + COINCIDENCE_TOLERANCE @@ -215,8 +221,8 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): if x > x_max - COINCIDENCE_TOLERANCE: i_x = mesh["Nx"] # - y_min = mesh_module.get_y(0, mesh, mcdc, data) - y_max = mesh_module.get_y(mesh["Ny"], mesh, mcdc, data) + y_min = mesh_module.get_y(0, mesh, simulation, data) + y_max = mesh_module.get_y(mesh["Ny"], mesh, simulation, data) if uy > 0.0: if ( y_final < y_min + COINCIDENCE_TOLERANCE @@ -234,8 +240,8 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): if y > y_max - COINCIDENCE_TOLERANCE: i_y = mesh["Ny"] # - z_min = mesh_module.get_z(0, mesh, mcdc, data) - z_max = mesh_module.get_z(mesh["Nz"], mesh, mcdc, data) + z_min = mesh_module.get_z(0, mesh, simulation, data) + z_max = mesh_module.get_z(mesh["Nz"], mesh, simulation, data) if uz > 0.0: if ( z_final < z_min + COINCIDENCE_TOLERANCE @@ -290,17 +296,17 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): axis_crossed = AXIS_T if mesh_tally: - mesh = mcdc["meshes"][tally["spatial_filter_ID"]] + mesh = simulation["meshes"][tally["spatial_filter_ID"]] # x-direction if ux == 0.0: dx = INF else: if ux > 0.0: - x_next = mesh_module.get_x(i_x + 1, mesh, mcdc, data) + x_next = mesh_module.get_x(i_x + 1, mesh, simulation, data) x_next = min(x_next, x_final) else: - x_next = mesh_module.get_x(i_x, mesh, mcdc, data) + x_next = mesh_module.get_x(i_x, mesh, simulation, data) x_next = max(x_next, x_final) dx = (x_next - x) / ux if dx <= distance_scored: @@ -312,10 +318,10 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): dy = INF else: if uy > 0.0: - y_next = mesh_module.get_y(i_y + 1, mesh, mcdc, data) + y_next = mesh_module.get_y(i_y + 1, mesh, simulation, data) y_next = min(y_next, y_final) else: - y_next = mesh_module.get_y(i_y, mesh, mcdc, data) + y_next = mesh_module.get_y(i_y, mesh, simulation, data) y_next = max(y_next, y_final) dy = (y_next - y) / uy if dy <= distance_scored: @@ -327,10 +333,10 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): dz = INF else: if uz > 0.0: - z_next = mesh_module.get_z(i_z + 1, mesh, mcdc, data) + z_next = mesh_module.get_z(i_z + 1, mesh, simulation, data) z_next = min(z_next, z_final) else: - z_next = mesh_module.get_z(i_z, mesh, mcdc, data) + z_next = mesh_module.get_z(i_z, mesh, simulation, data) z_next = max(z_next, z_final) dz = (z_next - z) / uz if dz <= distance_scored: @@ -348,21 +354,21 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): if score_type == SCORE_FLUX: score = flux elif score_type == SCORE_DENSITY: - speed = physics.particle_speed(particle_container, mcdc, data) + speed = physics.particle_speed(particle_container, simulation, data) score = flux / speed elif score_type == SCORE_COLLISION: score = flux * physics.macro_xs( - NEUTRON_REACTION_TOTAL, particle_container, mcdc, data + NEUTRON_REACTION_TOTAL, particle_container, simulation, data ) elif score_type == SCORE_CAPTURE: score = flux * physics.macro_xs( - NEUTRON_REACTION_CAPTURE, particle_container, mcdc, data + NEUTRON_REACTION_CAPTURE, particle_container, simulation, data ) elif score_type == SCORE_FISSION: score = flux * physics.macro_xs( - NEUTRON_REACTION_FISSION, particle_container, mcdc, data + NEUTRON_REACTION_FISSION, particle_container, simulation, data ) - atomic_add(data, idx_base + i_score, score) + util.atomic_add(data, idx_base + i_score, score) # Accumulate distance swept distance_swept += distance_scored @@ -381,7 +387,7 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): if i_time == tally_base["time_length"] - 1: return elif mesh_tally: - mesh = mcdc["meshes"][tally["spatial_filter_ID"]] + mesh = simulation["meshes"][tally["spatial_filter_ID"]] if axis_crossed == AXIS_X: if ux > 0.0: i_x += 1 @@ -423,33 +429,33 @@ def tracklength_tally(particle_container, distance, tally, mcdc, data): @njit -def eigenvalue_tally(particle_container, distance, mcdc, data): +def eigenvalue_tally(particle_container, distance, simulation, data): particle = particle_container[0] flux = distance * particle["w"] # Get nu-fission nuSigmaF = physics.neutron_production_xs( - NEUTRON_REACTION_FISSION, particle_container, mcdc, data + NEUTRON_REACTION_FISSION, particle_container, simulation, data ) # Fission production (needed even during inactive cycle) - atomic_add(mcdc["eigenvalue_tally_nuSigmaF"], 0, flux * nuSigmaF) + util.atomic_add(simulation["eigenvalue_tally_nuSigmaF"], 0, flux * nuSigmaF) # Done, if inactive - if not mcdc["cycle_active"]: + if not simulation["cycle_active"]: return # ================================================================================== # Neutron density # ================================================================================== - v = physics.particle_speed(particle_container, mcdc, data) + v = physics.particle_speed(particle_container, simulation, data) n_density = flux / v - atomic_add(mcdc["eigenvalue_tally_n"], 0, n_density) + util.atomic_add(simulation["eigenvalue_tally_n"], 0, n_density) # Maximum neutron density - if mcdc["n_max"] < n_density: - mcdc["n_max"] = n_density + if simulation["n_max"] < n_density: + simulation["n_max"] = n_density # ================================================================================== # TODO: Delayed neutron precursor density @@ -457,7 +463,7 @@ def eigenvalue_tally(particle_container, distance, mcdc, data): return # Get the decay-wighted multiplicity total = 0.0 - if mcdc["settings"]["multigroup_mode"]: + if simulation["settings"]["neutron_multigroup_mode"]: g = particle["g"] for j in range(J): nu_d = mcdc_get.material.mgxs_nu_d(g, j, material, data) @@ -467,7 +473,7 @@ def eigenvalue_tally(particle_container, distance, mcdc, data): E = P["E"] for i in range(material["N_nuclide"]): ID_nuclide = material["nuclide_IDs"][i] - nuclide = mcdc["nuclides"][ID_nuclide] + nuclide = simulation["nuclides"][ID_nuclide] if not nuclide["fissionable"]: continue for j in range(J): @@ -475,10 +481,12 @@ def eigenvalue_tally(particle_container, distance, mcdc, data): decay = nuclide["ce_decay"][j] total += nu_d / decay - SigmaF = physics.macro_xs(NEUTRON_REACTION_FISSION, particle_container, mcdc, data) - C_density = flux * total * SigmaF / mcdc["k_eff"] - atomic_add(mcdc["eigenvalue_tally_C"], 0, C_density) + SigmaF = physics.macro_xs( + NEUTRON_REACTION_FISSION, particle_container, simulation, data + ) + C_density = flux * total * SigmaF / simulation["k_eff"] + util.atomic_add(simulation["eigenvalue_tally_C"], 0, C_density) # Maximum precursor density - if mcdc["C_max"] < C_density: - mcdc["C_max"] = C_density + if simulation["C_max"] < C_density: + simulation["C_max"] = C_density diff --git a/mcdc/transport/technique.py b/mcdc/transport/technique.py index f8b15c05..fe629af3 100644 --- a/mcdc/transport/technique.py +++ b/mcdc/transport/technique.py @@ -9,6 +9,7 @@ import mcdc.transport.particle as particle_module import mcdc.transport.particle_bank as particle_bank_module import mcdc.transport.rng as rng +import mcdc.transport.util as util # ====================================================================================== # Weight Roulette @@ -16,10 +17,10 @@ @njit -def weight_roulette(particle_container, mcdc): +def weight_roulette(particle_container, simulation): particle = particle_container[0] - if particle["w"] < mcdc["weight_roulette"]["weight_threshold"]: - w_target = mcdc["weight_roulette"]["weight_target"] + if particle["w"] < simulation["weight_roulette"]["weight_threshold"]: + w_target = simulation["weight_roulette"]["weight_target"] survival_probability = particle["w"] / w_target if rng.lcg(particle_container) < survival_probability: particle["w"] = w_target @@ -33,15 +34,15 @@ def weight_roulette(particle_container, mcdc): @njit -def population_control(mcdc): +def population_control(simulation): """Uniform Splitting-Roulette technique""" - bank_census = mcdc["bank_census"] - M = mcdc["settings"]["N_particle"] - bank_source = mcdc["bank_source"] + bank_census = simulation["bank_census"] + M = simulation["settings"]["N_particle"] + bank_source = simulation["bank_source"] # Scan the bank - idx_start, N_local, N = particle_bank_module.bank_scanning(bank_census, mcdc) + idx_start, N_local, N = particle_bank_module.bank_scanning(bank_census, simulation) idx_end = idx_start + N_local # Abort if census bank is empty @@ -54,29 +55,29 @@ def population_control(mcdc): # Splitting Number sn = 1.0 / ws - P_rec_arr = np.zeros(1, type_.particle_data) + P_rec_arr = util.local_array(1, type_.particle_data) P_rec = P_rec_arr[0] # Perform split-roulette to all particles in local bank particle_bank_module.set_bank_size(bank_source, 0) for idx in range(N_local): # Weight of the surviving particles - w = bank_census["particles"][idx]["w"] + w = bank_census["particle_data"][idx]["w"] w_survive = w * ws # Determine number of guaranteed splits N_split = math.floor(sn) # Survive the russian roulette? - xi = rng.lcg(bank_census["particles"][idx : idx + 1]) + xi = rng.lcg(bank_census["particle_data"][idx : idx + 1]) if xi < sn - N_split: N_split += 1 # Split the particle for i in range(N_split): particle_module.copy_as_child( - P_rec_arr, bank_census["particles"][idx : idx + 1] + P_rec_arr, bank_census["particle_data"][idx : idx + 1] ) # Set weight P_rec["w"] = w_survive - particle_bank_module.bank_source_particle(P_rec_arr, mcdc) + particle_bank_module.bank_source_particle(P_rec_arr, simulation) diff --git a/mcdc/transport/util.py b/mcdc/transport/util.py index a82609ef..4627841d 100644 --- a/mcdc/transport/util.py +++ b/mcdc/transport/util.py @@ -1,13 +1,12 @@ import math +import numpy as np from numba import njit from typing import Sequence @njit -def find_bin( - value: float, grid: Sequence[float], epsilon: float = 0.0, go_lower: bool = True -) -> int: +def find_bin_with_rules(value, grid, epsilon, go_lower): """ Return the bin index i for which grid[i] <= value < grid[i+1], with optional epsilon tolerance and tie-breaking toward the lower/upper bin. @@ -18,10 +17,10 @@ def find_bin( Query point. grid : Sequence[float] Monotonically increasing bin edges of length N_grid = N_bin + 1. - epsilon : float, optional (default: 0.0) + epsilon : float Tolerance to treat values as being exactly on a grid edge if |value - grid[k]| <= epsilon. - go_lower : bool, optional (default: True) + go_lower : bool Tie-breaking rule when value is at/within epsilon of a grid edge: - True -> tie to the lower/left bin - False -> tie to the upper/right bin @@ -96,8 +95,16 @@ def find_bin( @njit -def atomic_add(array, idx, value): - array[idx] += value +def find_bin(value, grid): + tolerance = 0.0 + go_lower = True + return find_bin_with_rules(value, grid, tolerance, go_lower) + + +@njit +def find_bin_with_tolerance(value, grid, tolerance): + go_lower = True + return find_bin_with_rules(value, grid, tolerance, go_lower) # ====================================================================================== @@ -123,3 +130,23 @@ def log_interpolation(x, x1, x2, y1, y2): ly = ly1 + m * (math.log(x) - lx1) return math.exp(ly) + + +# ====================================================================================== +# Framework utilities +# ====================================================================================== + + +@njit +def atomic_add(array, idx, value): + array[idx] += value + + +@njit +def local_array(shape, dtype): + return np.zeros(shape, dtype=dtype) + + +@njit +def access_simulation(program): + return program diff --git a/mcdc/visualize.py b/mcdc/visualize.py index eeab164f..25640041 100644 --- a/mcdc/visualize.py +++ b/mcdc/visualize.py @@ -45,8 +45,8 @@ def visualize( global _visualize_cache if _visualize_cache is None: _visualize_cache = preparation() - mcdc_container, data = _visualize_cache - mcdc = mcdc_container[0] + simulation_container, data = _visualize_cache + simulation = simulation_container[0] # ================================================================================== # Numba-compiled functions @@ -64,7 +64,7 @@ def _compute_material_row( reference_val, time_val, particle_arr, - mcdc, + simulation, data, ): """ @@ -88,7 +88,7 @@ def _compute_material_row( Time value for the visualization particle_arr : np.ndarray Particle array of size (1,) used for particle lookup. - mcdc : structured array + simulation : structured array MCDC simulation data data : structured array Additional simulation data @@ -135,7 +135,7 @@ def _compute_material_row( particle["cell_ID"] = -1 particle["material_ID"] = -1 - if locate_particle(particle_arr, mcdc, data): + if locate_particle(particle_arr, simulation, data): row_materials[j] = particle["material_ID"] else: row_materials[j] = -1 @@ -152,7 +152,7 @@ def _compute_material_row( colors = new_colors else: colors = {} - for i in range(len(mcdc["materials"])): + for i in range(len(simulation["materials"])): colors[i] = plt.cm.Set1(i)[:-1] WHITE = mpl_colors.to_rgb("white") @@ -222,7 +222,7 @@ def _compute_material_row( reference, t, particle_arr, - mcdc, + simulation, data, ) material_ids[i, :] = row_materials