diff --git a/.gitignore b/.gitignore index 9f7bc907..ef42397c 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ *.dat *.mcv out.* +log.* compiled_composite* composite_kernel_* @@ -164,4 +165,4 @@ dmypy.json cython_debug/ # PyCharm -.idea/ \ No newline at end of file +.idea/ diff --git a/acc/_numba.py b/acc/_numba.py index f9313951..1d4c6631 100644 --- a/acc/_numba.py +++ b/acc/_numba.py @@ -1,3 +1,4 @@ +import os, tempfile import numpy as np from numba.core import config if config.ENABLE_CUDASIM: @@ -6,3 +7,38 @@ def xoroshiro128p_uniform_float32(rng_states, threadindex): from numba.cuda.random import xoroshiro128p_type else: from numba.cuda.random import xoroshiro128p_uniform_float32, xoroshiro128p_type + +class Coder: + + def __init__(self, workdir=None): + self.workdir = workdir = workdir or os.path.abspath(".mcvine.acc.coder") + if not os.path.exists(workdir): + os.makedirs(workdir) + self.modules = dict() + + def getModule(self, type, N): + container = coder.createDir(type) + modulepath = os.path.join(container, f'compiled_{N}.py') + key = type, N + if os.path.exists(modulepath) and key not in self.modules: + self.modules[key] = modulepath + return modulepath + + def createUniqueDir(self, prefix): + return tempfile.mkdtemp(prefix=prefix, dir=self.workdir) + + def createDir(self, name): + wd = os.path.join(self.workdir, name) + if not os.path.exists(wd): + os.makedirs(wd) + return wd + + @classmethod + def unrollLoop(cls, N, indent=4*' ', before_loop=None, in_loop=None, after_loop=None): + lines=[indent+line for line in before_loop or []] + for i in range(N): + lines+=[indent+line.format(i=i) for line in in_loop or []] + lines+=[indent+line for line in after_loop or []] + return lines + +coder = Coder() diff --git a/acc/components/samples/composite.py b/acc/components/samples/composite.py new file mode 100644 index 00000000..5ff2a4be --- /dev/null +++ b/acc/components/samples/composite.py @@ -0,0 +1,48 @@ +import os +import numpy as np, numba +from numba import cuda, void, int64 +from numba.cuda.random import xoroshiro128p_uniform_float32, xoroshiro128p_type +from ...config import get_numba_floattype +NB_FLOAT = get_numba_floattype() + +from .SampleBase import SampleBase + +category = 'samples' + +def sampleassembly_from_xml(name, samplexml): + """factory method to create sample assembly component. + can only be used in nonacc instrument script + """ + klass = SampleAssemblyFromXml(samplexml) + return klass(name) + +def SampleAssemblyFromXml(samplexml): + from . import loadScattererComposite + composite = loadScattererComposite(samplexml) + return factory(composite) + +def factory(composite): + from ...scatterers import scatter_func_factory + methods = scatter_func_factory.render(composite) + scatter = methods['scatter'] + class Composite(SampleBase): + + def __init__(self, name): + self.name = name + self.propagate_params = () + + # Aim neutrons toward the sample to cause JIT compilation. + import mcni + neutrons = mcni.neutron_buffer(1) + neutrons[0] = mcni.neutron(r=(0, 0, -1), v=(0, 0, 1), prob=1, time=0) + self.process(neutrons) + + @cuda.jit( + void(int64, xoroshiro128p_type[:], NB_FLOAT[:]), + device=True, inline=True, + ) + def propagate(threadindex, rng_states, neutron): + return + Composite.propagate = scatter + Composite.register_propagate_method(scatter) + return Composite diff --git a/acc/geometry/arrow_intersect.py b/acc/geometry/arrow_intersect.py index 830a933e..c6860082 100644 --- a/acc/geometry/arrow_intersect.py +++ b/acc/geometry/arrow_intersect.py @@ -26,6 +26,18 @@ def render(self, shape): return shape.identify(self) def onUnion(self, u): + nelements = len(u.shapes) + if nelements == 1: + return u.shapes[0].identify(self) + elif nelements == 2: + return self.onUnion2(u) + elif nelements == 3: + from .composite_3 import createRayTracingMethods_NonOverlappingShapes as createMethods + return createMethods(u.shapes)['intersect_all'] + else: + raise NotImplementedError(f"locate for union of {nelements} elements") + + def onUnion2(self, u): locate1 = self.locate_func_factory.onUnion(u) s1, s2 = u.shapes f1 = s1.identify(self) @@ -85,6 +97,22 @@ def intersectBlock(x, y, z, vx, vy, vz, ts): return intersectBlock + def onTranslation(self, t): + x0,y0,z0 = [_/units.length.meter for _ in t.vector] + body = t.body + intersect_body = self.render(body) + @cuda.jit(device=True, inline=True) + def intersectTranslation(x,y,z, vx,vy,vz, ts): + return intersect_body(x-x0,y-y0,z-z0, vx,vy,vz, ts) + return intersectTranslation + + def onRotation(self, r): + euler_angles = [_/units.angle.rad for _ in r.euler_angles] + # hack + if not np.allclose(euler_angles, [0,0,0]): + raise NotImplementedError + return self.render(r.body) + def onDifference(self, s): locate1 = self.locate_func_factory.onDifference(s) f1 = s.op1.identify(self) @@ -150,8 +178,13 @@ def remove_item(idx, l, N): @cuda.jit(device=True) def insert_into_sorted_list(d, l, N): - 'insert data "d" into existing sorted array (low to high) of length N' - if N>=len(l): return N + '''insert data "d" into existing sorted array (low to high) of length N + and keep the lower number when there are too many elements + ''' + if N>=len(l): + if l[N-1] <= d: + return N + N = N - 1 if N==0: l[0] = d return 1 diff --git a/acc/geometry/composite.py b/acc/geometry/composite.py new file mode 100644 index 00000000..0880f2fc --- /dev/null +++ b/acc/geometry/composite.py @@ -0,0 +1,309 @@ +import os, numpy as np, numba +from numba import cuda +from mcvine.acc.geometry.arrow_intersect import max_intersections +from mcvine.acc import test +from .._numba import coder + +def get_find_1st_hit(shapes): + N = len(shapes) + mod = _importModule(N) + methods = mod.createRayTracingMethods_NonOverlappingShapes(shapes) + return _make_find_1st_hit(**methods) + +def get_union_locate(shapes): + N = len(shapes) + mod = _importModule(N) + return mod.createUnionLocateMethod(shapes) + +def _make_find_1st_hit(forward_intersect_all, is_onborder, find_shape_containing_point, **kwds): + @cuda.jit(device=True) + def _find_1st_hit(x,y,z, vx,vy,vz, ts): + nIntersections = forward_intersect_all(x,y,z, vx,vy,vz, ts) + # we have two cases + # case 1 + # shape1 | vacuum | shape2 + # start-|---------|---> + # case 2 + # vacuum | shape1 + # start-|-----------> + # case 1: there will be odd number of intersections + # case 2: there will be even number of intersections + # we just need to determine which one of the above two cases is true, + # case1: + if nIntersections % 2 == 1: + return find_shape_containing_point( + x + ts[0]/2.*vx, + y + ts[0]/2.*vy, + z + ts[0]/2.*vz, + ) + else : + # case2: + # no intersection + if nIntersections==0: + return -1 + # at least two + midt = (ts[0]+ts[1])/2. + ret = find_shape_containing_point( + x + midt*vx, + y + midt*vy, + z + midt*vz, + ) + # let us determine if the start is on border + isonborder = is_onborder(x,y,z) + # If start is not on border of any shape, it would be easier. + if (not isonborder): return ret + + # on border. that is a bit more difficult. + # we need to go over all intersection pairs, and find the first pair + # whose midlle point is insde a shape. That shape containing + # the middle point is the target. + previous = 0.0 + for point_index in range(nIntersections): + now = ts[point_index] + midt = (previous+now)/2 + ret = find_shape_containing_point( + x + midt*vx, + y + midt*vy, + z + midt*vz, + ) + if ret>=0: return ret + previous = now + return -1 + if test.USE_CUDASIM: + @cuda.jit(device=True) + def find_1st_hit(x,y,z, vx,vy,vz): + ts = np.zeros(max_intersections, dtype=float) + return _find_1st_hit(x,y,z, vx,vy,vz, ts) + else: + @cuda.jit(device=True) + def find_1st_hit(x,y,z, vx,vy,vz): + ts = cuda.local.array(max_intersections, dtype=numba.float64) + return _find_1st_hit(x,y,z, vx,vy,vz, ts) + return find_1st_hit + + +def _importModule(N): + mod = _makeModule(N) + import imp + return imp.load_source(f'geometry_composite_{N}', mod) + +def _makeModule(N, overwrite=False): + "make cuda device methods for composite with N elements" + from .._numba import coder + modulepath = coder.getModule("composite_shape", N) + if os.path.exists(modulepath) and not overwrite: + return modulepath + indent = 4*' ' + imports = """import os, numpy as np, numba +from numba import cuda +from mcvine.acc._numba import xoroshiro128p_uniform_float32 +from mcvine.acc import test +from mcvine.acc.geometry import arrow_intersect +from mcvine.acc.geometry.location import inside, onborder, outside +from mcvine.acc.geometry.arrow_intersect import max_intersections, insert_into_sorted_list +""".splitlines() + createRTMethods = _Coder_createRTMethods(N, indent)() + createUnionLocateMethod = _Coder_createUnionLocateMethod(N, indent)() + lines = imports + [''] + createRTMethods + createUnionLocateMethod + with open(modulepath, 'wt') as ostream: + ostream.write("\n".join(lines)) + return modulepath + +class _Coder_createUnionLocateMethod: + + def __init__(self, N, indent=4*' '): + self.N = N + self.indent = indent + return + + def __call__(self): + N, indent = self.N, self.indent + header = f"""assert len(shapes)=={N} +locates = [ + arrow_intersect.locate_func_factory.render(shape) + for shape in shapes +] +""".splitlines() + locate_loop = coder.unrollLoop( + N = N, + indent = '', + in_loop = ["locate_{i} = locates[{i}]"], + ) + body = header + locate_loop + [''] + self.locate() + add_indent = lambda lines, n: [indent*n+l for l in lines] + return ( + ["def createUnionLocateMethod(shapes):"] + + add_indent(body, 1) + + [indent+'return locate'] + ) + + def locate(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def locate(x,y,z):", + ] + assign_loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = ["loc{i} = locate_{i}(x,y,z)"], + ) + inside_loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = ["if loc{i} == inside: return inside"], + ) + onborder_loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = ["if loc{i} == onborder: return onborder"], + after_loop = ['return outside'] + ) + return header + assign_loop + inside_loop + onborder_loop + + +class _Coder_createRTMethods: + + def __init__(self, N, indent=4*' '): + self.N = N + self.indent = indent + return + + def __call__(self): + N, indent = self.N, self.indent + header = f"""assert len(shapes)=={N} +funcs_list = [ + ( + arrow_intersect.locate_func_factory.render(shape), + arrow_intersect.arrow_intersect_func_factory.render(shape), + ) + for shape in shapes +] + """.splitlines() + funcs_loop = coder.unrollLoop( + N = N, + indent = '', + in_loop = ["locate_{i}, intersect_{i} = funcs_list[{i}]"], + ) + end = """ +if test.USE_CUDASIM: + @cuda.jit(device=True) + def intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = np.zeros(max_intersections, dtype=float) + return _intersect_all(x,y,z, vx,vy,vz, ts, ts_) + @cuda.jit(device=True) + def forward_intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = np.zeros(max_intersections, dtype=float) + return _forward_intersect_all(x,y,z, vx,vy,vz, ts, ts_) +else: + @cuda.jit(device=True) + def intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = cuda.local.array(max_intersections, dtype=numba.float64) + return _intersect_all(x,y,z, vx,vy,vz, ts, ts_) + @cuda.jit(device=True) + def forward_intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = cuda.local.array(max_intersections, dtype=numba.float64) + return _forward_intersect_all(x,y,z, vx,vy,vz, ts, ts_) +return dict( + intersect_all = intersect_all, + forward_intersect_all = forward_intersect_all, + find_shape_containing_point = find_shape_containing_point, + is_onborder = is_onborder, +) + """.splitlines() + body = ( + header + funcs_loop + [''] + + self._intersect_all() + + self._forward_intersect_all() + + self.find_shape_containing_point() + + self.is_onborder() + + end + ) + add_indent = lambda lines, n: [indent*n+l for l in lines] + return ["def createRayTracingMethods_NonOverlappingShapes(shapes):"] + add_indent(body, 1) + + + def _intersect_all(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def _intersect_all(x,y,z, vx,vy,vz, ts, ts_):", + ] + loop = coder.unrollLoop( + N = N, + indent = indent, + before_loop = ["N=0"], + in_loop = [ + "N_ = intersect_{i}(x,y,z, vx,vy,vz, ts_)", + "for i in range(N_):", + indent + "N = insert_into_sorted_list(ts_[i], ts, N)", + ], + after_loop = ["return N"] + ) + return header + loop + + def _forward_intersect_all(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def _forward_intersect_all(x,y,z, vx,vy,vz, ts, ts_):", + ] + loop = coder.unrollLoop( + N = N, + indent = indent, + before_loop = ["N=0"], + in_loop = [ + "N_ = intersect_{i}(x,y,z, vx,vy,vz, ts_)", + "for i in range(N_):", + indent + "t = ts_[i]", + indent + "if t>0:", + indent*2 + "N = insert_into_sorted_list(t, ts, N)", + ], + after_loop = ["return N"] + ) + return header + loop + + def find_shape_containing_point(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def find_shape_containing_point(x,y,z):", + ] + loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = [ + "if locate_{i}(x,y,z) == inside:", + indent + "return {i}", + ], + after_loop = ["return -1"] + ) + return header + loop + + def is_onborder(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def is_onborder(x,y,z):", + ] + loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = [ + "if locate_{i}(x,y,z) == onborder:", + indent + "return True", + ], + after_loop = ["return False"] + ) + return header + loop + +module_code_template = """ +import os, numpy as np, numba +from numba import cuda +from mcvine.acc._numba import xoroshiro128p_uniform_float32 +from mcvine.acc import test +from mcvine.acc.geometry import arrow_intersect +from mcvine.acc.geometry.location import inside, onborder, outside +from mcvine.acc.geometry.arrow_intersect import max_intersections, insert_into_sorted_list + +""" diff --git a/acc/geometry/composite_3.py b/acc/geometry/composite_3.py new file mode 100644 index 00000000..02a80c63 --- /dev/null +++ b/acc/geometry/composite_3.py @@ -0,0 +1,129 @@ +# Example code that would be generated automatically + +import os, numpy as np, numba +from numba import cuda +from mcvine.acc._numba import xoroshiro128p_uniform_float32 +from mcvine.acc import test +from mcvine.acc.geometry import arrow_intersect +from mcvine.acc.geometry.location import inside, onborder, outside +from mcvine.acc.geometry.arrow_intersect import max_intersections, insert_into_sorted_list + +def createRayTracingMethods_NonOverlappingShapes(shapes): + "methods for a list of non-overlapping shapes" + assert len(shapes)==3 + funcs_list = [ + ( + arrow_intersect.locate_func_factory.render(shape), + arrow_intersect.arrow_intersect_func_factory.render(shape), + ) + for shape in shapes + ] + + locate_0, intersect_0 = funcs_list[0] + locate_1, intersect_1 = funcs_list[1] + locate_2, intersect_2 = funcs_list[2] + + @cuda.jit(device=True) + def _intersect_all(x,y,z, vx,vy,vz, ts, ts_): + N = 0 + N_ = intersect_0(x,y,z, vx,vy,vz, ts_) + for i in range(N_): + N = insert_into_sorted_list(ts_[i], ts, N) + N_ = intersect_1(x,y,z, vx,vy,vz, ts_) + for i in range(N_): + N = insert_into_sorted_list(ts_[i], ts, N) + N_ = intersect_2(x,y,z, vx,vy,vz, ts_) + for i in range(N_): + N = insert_into_sorted_list(ts_[i], ts, N) + return N + + @cuda.jit(device=True) + def _forward_intersect_all(x,y,z, vx,vy,vz, ts, ts_): + N = 0 + N_ = intersect_0(x,y,z, vx,vy,vz, ts_) + for i in range(N_): + t = ts_[i] + if t>0: + N = insert_into_sorted_list(t, ts, N) + N_ = intersect_1(x,y,z, vx,vy,vz, ts_) + for i in range(N_): + t = ts_[i] + if t>0: + N = insert_into_sorted_list(t, ts, N) + N_ = intersect_2(x,y,z, vx,vy,vz, ts_) + for i in range(N_): + t = ts_[i] + if t>0: + N = insert_into_sorted_list(t, ts, N) + return N + + @cuda.jit(device=True) + def find_shape_containing_point(x,y,z): + if locate_0(x,y,z) == inside: + return 0 + if locate_1(x,y,z) == inside: + return 1 + if locate_2(x,y,z) == inside: + return 2 + return -1 + + @cuda.jit(device=True) + def is_onborder(x,y,z): + "check if the point is on the border of any of the shapes" + if locate_0(x,y,z) == onborder: + return True + if locate_1(x,y,z) == onborder: + return True + if locate_2(x,y,z) == onborder: + return True + return False + + if test.USE_CUDASIM: + @cuda.jit(device=True) + def intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = np.zeros(max_intersections, dtype=float) + return _intersect_all(x,y,z, vx,vy,vz, ts, ts_) + @cuda.jit(device=True) + def forward_intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = np.zeros(max_intersections, dtype=float) + return _forward_intersect_all(x,y,z, vx,vy,vz, ts, ts_) + else: + @cuda.jit(device=True) + def intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = cuda.local.array(max_intersections, dtype=numba.float64) + return _intersect_all(x,y,z, vx,vy,vz, ts, ts_) + @cuda.jit(device=True) + def forward_intersect_all(x,y,z, vx,vy,vz, ts): + ts_ = cuda.local.array(max_intersections, dtype=numba.float64) + return _forward_intersect_all(x,y,z, vx,vy,vz, ts, ts_) + return dict( + intersect_all = intersect_all, + forward_intersect_all = forward_intersect_all, + find_shape_containing_point = find_shape_containing_point, + is_onborder = is_onborder, + ) + +def createUnionLocateMethod(shapes): + assert len(shapes)==3 + locates = [ + arrow_intersect.locate_func_factory.render(shape) + for shape in shapes + ] + locate_0 = locates[0] + locate_1 = locates[1] + locate_2 = locates[2] + + @cuda.jit(device=True) + def locate(x,y,z): + loc0 = locate_0(x,y,z) + loc1 = locate_1(x,y,z) + loc2 = locate_2(x,y,z) + if loc0 == inside: return inside + if loc1 == inside: return inside + if loc2 == inside: return inside + if loc0 == onborder: return onborder + if loc1 == onborder: return onborder + if loc2 == onborder: return onborder + return outside + + return locate diff --git a/acc/geometry/locate.py b/acc/geometry/locate.py index 9fd489a9..0d757ab8 100644 --- a/acc/geometry/locate.py +++ b/acc/geometry/locate.py @@ -1,3 +1,4 @@ +import numpy as np from numba import cuda from mcni import units from . import epsilon, location @@ -11,6 +12,18 @@ def render(self, shape): return shape.identify(self) def onUnion(self, u): + nelements = len(u.shapes) + if nelements == 1: + return u.shapes[0].identify(self) + elif nelements == 2: + return self.onUnion2(u) + elif nelements == 3: + from .composite_3 import createUnionLocateMethod + return createUnionLocateMethod(u.shapes) + else: + raise NotImplementedError(f"locate for union of {nelements} elements") + + def onUnion2(self, u): s1, s2 = u.shapes f1 = s1.identify(self) f2 = s2.identify(self) @@ -42,9 +55,24 @@ def onBlock(self, b): @cuda.jit(device=True, inline=True) def locateWrtBlock(x, y, z): return cu_device_locate_wrt_box(x, y, z, W, H, D) - return locateWrtBlock + def onTranslation(self, t): + x0,y0,z0 = [_/units.length.meter for _ in t.vector] + body = t.body + locate_in_body = self.render(body) + @cuda.jit(device=True, inline=True) + def locateWrtTranslation(x,y,z): + return locate_in_body(x-x0, y-y0, z-z0) + return locateWrtTranslation + + def onRotation(self, r): + euler_angles = [_/units.angle.rad for _ in r.euler_angles] + # hack + if not np.allclose(euler_angles, [0,0,0]): + raise NotImplementedError + return self.render(r.body) + def onDifference(self, s): f1 = s.op1.identify(self) f2 = s.op2.identify(self) @@ -52,7 +80,6 @@ def onDifference(self, s): @cuda.jit(device=True, inline=True) def locateWrtDifference(x, y, z): return cu_device_locate_wrt_difference(x, y, z, f1, f2) - return locateWrtDifference def onIntersection(self, u): diff --git a/acc/geometry/propagation.py b/acc/geometry/propagation.py index 3df424a4..acd4a199 100644 --- a/acc/geometry/propagation.py +++ b/acc/geometry/propagation.py @@ -4,6 +4,7 @@ from .arrow_intersect import max_intersections from .arrow_intersect import inside, outside, onborder from ..neutron import prop_dt_inplace +from ..vec3 import distance as v3_dist, length as v3_length from .. import test def makePropagateMethods(intersect, locate): @@ -190,6 +191,34 @@ def _tof_before_first_exit(neutron, ts): continue return t + @cuda.jit(device=True) + def _forward_distance_in_shape(neutron, end, ts): + "distance to travel in shape to get to `end`. `end` must be on the path of the neutron" + start = neutron[:3] + vv = neutron[3:6] + v = v3_length(vv) + nintersect = forward_intersect(start[0], start[1], start[2], vv[0], vv[1], vv[2], ts) + length = v3_dist(start, end) + tofmax = length/v; + prev = 0; length = 0 + for i in range(nintersect): + tof = ts[i] + if (tof > tofmax): + tof = tofmax + middle = (tof+prev)/2. + # middle point p = start + vv*middle; + x = start[0] + vv[0]*middle + y = start[1] + vv[1]*middle + z = start[2] + vv[2]*middle + # if middle point is inside, count this segment + if locate(x,y,z) == inside: + length += (tof-prev) * v; + if tof > tofmax: + break + prev = tof + continue + return length + if test.USE_CUDASIM: @cuda.jit(device=True, inline=True) @@ -212,6 +241,10 @@ def propagate_to_next_incident_surface(neutron): def propagate_to_next_exiting_surface(neutron): ts = np.zeros(max_intersections, dtype=float) return _propagate_to_next_exiting_surface(neutron, ts) + @cuda.jit(device=True, inline=True) + def forward_distance_in_shape(neutron, end): + ts = np.zeros(max_intersections, dtype=float) + return _forward_distance_in_shape(neutron, end, ts) else: @cuda.jit(device=True, inline=True) def is_exiting(neutron): @@ -233,6 +266,10 @@ def propagate_to_next_incident_surface(neutron): def propagate_to_next_exiting_surface(neutron): ts = cuda.local.array(max_intersections, dtype=numba.float64) return _propagate_to_next_exiting_surface(neutron, ts) + @cuda.jit(device=True, inline=True) + def forward_distance_in_shape(neutron, end): + ts = cuda.local.array(max_intersections, dtype=numba.float64) + return _forward_distance_in_shape(neutron, end, ts) return dict( is_exiting = is_exiting, @@ -241,4 +278,5 @@ def propagate_to_next_exiting_surface(neutron): propagate_out = propagate_out, propagate_to_next_incident_surface = propagate_to_next_incident_surface, propagate_to_next_exiting_surface = propagate_to_next_exiting_surface, + forward_distance_in_shape = forward_distance_in_shape, ) diff --git a/acc/kernels/composite.py b/acc/kernels/composite.py index fb01c009..d801da1f 100644 --- a/acc/kernels/composite.py +++ b/acc/kernels/composite.py @@ -9,119 +9,223 @@ def makeKernelMethods(composite): mod = makeKernelModule(composite) import imp m = imp.load_source('composite', mod) - return m.scatter, m.scattering_coeff, m.absorb, m.absorption_coeff + return m.createKernelMethods(composite) -_modules = {} def makeKernelModule(composite): "make cuda device methods for the composite kernel" - if composite in _modules: return _modules[composite] - import pickle as pkl, tempfile + from .._numba import coder + nkernels = len(composite.elements()) + modulepath = coder.getModule("composite_kernel", nkernels) + if os.path.exists(modulepath): + return modulepath # save composite to be loaded by the "compiled" module - coder_dir=os.path.join(os.curdir, '.mcvine.acc.coder') - if not os.path.exists(coder_dir): - os.makedirs(coder_dir) - tmpdir = tempfile.mkdtemp(prefix='composite_kernel_', dir=coder_dir) - pklpath = os.path.abspath(os.path.join(tmpdir, 'composite.pkl')) - pkl.dump(composite, open(pklpath, 'wb')) # make "compiled" composite module - nkernels = len(composite.elements()) + # + indent = 4*' ' + add_indent = lambda lines, n: [indent*n+l for l in lines] # 1. scatter element_scatter_method_defs = [ f'scatter_{ik}, scattering_coeff_{ik}, absorb_{ik}, absorption_coeff_{ik} = kernel_funcs_list[{ik}]' - for ik in range(nkernels)] - element_scatter_method_defs = '\n'.join(element_scatter_method_defs) - indent = 4*' ' - lines = _create_select_kernel_func_lines( - nkernels, - method='scatter', - args = 'threadindex, rng_states, neutron', - indent = indent - ) - if_clause_for_scatter_method = '\n'.join([indent + line for line in lines]) + for ik in range(nkernels) + ] + element_scatter_method_defs = '\n'.join(add_indent(element_scatter_method_defs, 1)) + scatter_method = '\n'.join(add_indent( + _create_scatter_method(nkernels, indent=indent), 1 + )) # 2. scattering_coeff - lines = [f'r += scattering_coeff_{i}(neutron)' for i in range(nkernels)] - add_scattering_coeff = '\n'.join(indent+l for l in lines) + scattering_coeff_method = '\n'.join(add_indent( + _create_scattering_coeff_method(nkernels, indent=4*' '), 1 + )) # 3. absorb - lines = _create_select_kernel_func_lines( - nkernels, - method='absorb', - args = 'threadindex, rng_states, neutron', - indent = indent - ) - if_clause_for_absorb_method = '\n'.join([indent + line for line in lines]) + absorb_method = '\n'.join(add_indent( + _create_absorb_method(nkernels, indent=indent), 1 + )) # 4. absorption_coeff - lines = [f'r += absorption_coeff_{i}(neutron)' for i in range(nkernels)] - add_absorption_coeff = '\n'.join(indent+l for l in lines) + abs_coeff_method = '\n'.join(add_indent( + _create_abs_coeff_method(nkernels, indent=4*' '), 1 + )) # all together content = template.format(**locals()) - modulepath = os.path.join(tmpdir, 'compiled_composite.py') open(modulepath, 'wt').write(content) - _modules[composite] = modulepath return modulepath +def _create_scattering_coeff_method(nkernels, indent=4*' '): + header = [ + '@cuda.jit(device=True)', + 'def scattering_coeff(neutron):' + ] + from .._numba import coder + loop = coder.unrollLoop( + N = nkernels, + indent = indent, + before_loop = ['r = 0.0'], + in_loop = ['r += scattering_coeff_{i}(neutron)'], + after_loop = ['return r*scale_scattering_coeff'] + ) + return header + loop + +def _create_scatter_method(nkernels, indent=4*' '): + header = [ + '@cuda.jit(device=True)', + 'def scatter(threadindex, rng_states, neutron):' + ] + loop = _create_select_kernel_func_lines( + nkernels, + method = 'scatter', + args = 'threadindex, rng_states, neutron', + indent = indent + ) + return header + loop + +def _create_abs_coeff_method(nkernels, indent=4*' '): + header = [ + '@cuda.jit(device=True)', + 'def absorption_coeff(neutron):' + ] + from .._numba import coder + loop = coder.unrollLoop( + N = nkernels, + indent = indent, + before_loop = ['r = 0.0'], + in_loop = ['r += absorption_coeff_{i}(neutron)'], + after_loop = ['return r/Nkernels'] + ) + return header + loop + +def _create_absorb_method(nkernels, indent=4*' '): + header = [ + '@cuda.jit(device=True)', + 'def absorb(threadindex, rng_states, neutron):' + ] + loop = _create_select_kernel_func_lines( + nkernels, + method = 'absorb', + args = 'threadindex, rng_states, neutron', + indent = indent + ) + return header + loop + def _create_select_kernel_func_lines(nkernels, method, args, indent=4*' '): - lines = [] - for i in range(nkernels-1): - lead = 'if' - if i>0: lead = 'elif' - lines.append(f'{lead} r < device_cumulative_weights[{i}]:') - lines.append(f'{indent}ret = {method}_{i}({args})') - lines.append(f'{indent}ikernel={i}') - lines.append(f'else:') - lines.append(f'{indent}ret = {method}_{nkernels-1}({args})') - lines.append(f'{indent}ikernel={nkernels-1}') - lines.append(f'neutron[-1]/=device_weights[ikernel]') - lines.append(f'return ret') - return lines - -template = """import os, pickle as pkl, numpy as np + from .._numba import coder + return coder.unrollLoop( + N = nkernels, + indent = indent, + before_loop = [ + 'r = xoroshiro128p_uniform_float32(rng_states, threadindex)' + ], + in_loop = [ + 'if r < device_cumulative_weights[{i}]:', + f' ret = {method}_'+'{i}'+f'({args})', + ' neutron[-1]/=device_weights[{i}]', + ' return ret' + ], + after_loop = [] + ) + +template = """ +import os, numpy as np from numba import cuda from mcvine.acc._numba import xoroshiro128p_uniform_float32 -pklpath = {pklpath!r} -composite = pkl.load(open(pklpath, 'rb')) -elements = composite.elements() -Nkernels = len(elements) -weights = [float(element.weight) for element in elements] -cumulative_weights = [weights[0]] -for w in weights[1:]: - cumulative_weights.append(w+cumulative_weights[-1]) -cumulative_weights = np.array(cumulative_weights) -device_cumulative_weights = cuda.to_device(cumulative_weights) -weights = np.array(weights) -device_weights = cuda.to_device(weights) -scale_scattering_coeff = 1./Nkernels if composite.average else 1 - -# create cuda functions for kernels -from mcvine.acc.kernels import scatter_func_factory -kernel_funcs_list = [] -for element in elements: - kernel_funcs = scatter_func_factory.render(element) - kernel_funcs_list.append(kernel_funcs) - continue - -# and assign names +def createKernelMethods(composite): + elements = composite.elements() + Nkernels = len(elements) + weights = [float(element.weight) for element in elements] + cumulative_weights = [weights[0]] + for w in weights[1:]: + cumulative_weights.append(w+cumulative_weights[-1]) + cumulative_weights = np.array(cumulative_weights) + device_cumulative_weights = cuda.to_device(cumulative_weights) + weights = np.array(weights) + device_weights = cuda.to_device(weights) + scale_scattering_coeff = 1./Nkernels if composite.average else 1 + + # create cuda functions for kernels + from mcvine.acc.kernels import scatter_func_factory + kernel_funcs_list = [] + for element in elements: + kernel_funcs = scatter_func_factory.render(element) + kernel_funcs_list.append(kernel_funcs) + continue + {element_scatter_method_defs} -@cuda.jit(device=True) -def scatter(threadindex, rng_states, neutron): - r = xoroshiro128p_uniform_float32(rng_states, threadindex) -{if_clause_for_scatter_method} - -@cuda.jit(device=True) -def scattering_coeff(neutron): - r = 0.0 -{add_scattering_coeff} - return r*scale_scattering_coeff - -@cuda.jit(device=True) -def absorb(threadindex, rng_states, neutron): - r = xoroshiro128p_uniform_float32(rng_states, threadindex) -{if_clause_for_absorb_method} - -@cuda.jit(device=True) -def absorption_coeff(neutron): - r = 0.0 -{add_absorption_coeff} - return r/Nkernels +{scatter_method} + +{absorb_method} + +{scattering_coeff_method} + +{abs_coeff_method} + return scatter, scattering_coeff, absorb, absorption_coeff +""" + +# Example code generated +""" +import os, numpy as np +from numba import cuda +from mcvine.acc._numba import xoroshiro128p_uniform_float32 + +def createKernelMethods(composite): + elements = composite.elements() + Nkernels = len(elements) + weights = [float(element.weight) for element in elements] + cumulative_weights = [weights[0]] + for w in weights[1:]: + cumulative_weights.append(w+cumulative_weights[-1]) + cumulative_weights = np.array(cumulative_weights) + device_cumulative_weights = cuda.to_device(cumulative_weights) + weights = np.array(weights) + device_weights = cuda.to_device(weights) + scale_scattering_coeff = 1./Nkernels if composite.average else 1 + + # create cuda functions for kernels + from mcvine.acc.kernels import scatter_func_factory + kernel_funcs_list = [] + for element in elements: + kernel_funcs = scatter_func_factory.render(element) + kernel_funcs_list.append(kernel_funcs) + continue + + scatter_0, scattering_coeff_0, absorb_0, absorption_coeff_0 = kernel_funcs_list[0] + scatter_1, scattering_coeff_1, absorb_1, absorption_coeff_1 = kernel_funcs_list[1] + + @cuda.jit(device=True) + def scatter(threadindex, rng_states, neutron): + r = xoroshiro128p_uniform_float32(rng_states, threadindex) + if r < device_cumulative_weights[0]: + ret = scatter_0(threadindex, rng_states, neutron) + neutron[-1]/=device_weights[0] + return ret + if r < device_cumulative_weights[1]: + ret = scatter_1(threadindex, rng_states, neutron) + neutron[-1]/=device_weights[1] + return ret + + @cuda.jit(device=True) + def absorb(threadindex, rng_states, neutron): + r = xoroshiro128p_uniform_float32(rng_states, threadindex) + if r < device_cumulative_weights[0]: + ret = absorb_0(threadindex, rng_states, neutron) + neutron[-1]/=device_weights[0] + return ret + if r < device_cumulative_weights[1]: + ret = absorb_1(threadindex, rng_states, neutron) + neutron[-1]/=device_weights[1] + return ret + + @cuda.jit(device=True) + def scattering_coeff(neutron): + r = 0.0 + r += scattering_coeff_0(neutron) + r += scattering_coeff_1(neutron) + return r*scale_scattering_coeff + + @cuda.jit(device=True) + def absorption_coeff(neutron): + r = 0.0 + r += absorption_coeff_0(neutron) + r += absorption_coeff_1(neutron) + return r/Nkernels + return scatter, scattering_coeff, absorb, absorption_coeff """ diff --git a/acc/scatterers/__init__.py b/acc/scatterers/__init__.py new file mode 100644 index 00000000..ec52fa13 --- /dev/null +++ b/acc/scatterers/__init__.py @@ -0,0 +1,27 @@ +import numba +from numba import cuda +from mcni import units + +class ScatterFuncFactory: + + def render(self, scatterer): + """returns cuda device function tuple (scatter, interact_path1) + """ + return scatterer.identify(self) + + def onCompositeScatterer(self, composite): + elements = composite.elements() + from .composite_scatterer import factory + return factory(composite) + + def onHomogeneousScatterer(self, hs): + from .homogeneous_scatterer import factory + shape = hs.shape() + kernel = hs.kernel() + mcweights = [float(w) for w in hs.mcweights] + packing_factor = float(hs.packing_factor) + methods = factory(shape, kernel, mcweights, packing_factor) + return methods + + +scatter_func_factory = ScatterFuncFactory() diff --git a/acc/scatterers/composite_3.py b/acc/scatterers/composite_3.py new file mode 100644 index 00000000..c4ee27d5 --- /dev/null +++ b/acc/scatterers/composite_3.py @@ -0,0 +1,55 @@ +# -*- python -*- +# +# Jiao Lin +# + +import numpy as np, numba +from numba import cuda +from numba.core import config + +from .interaction_types import none +from .. import test + +def createHelperMethodsForScatter(composite): + # elements + elements = composite.elements() + nelements = len(elements) + assert nelements == 3 + # methods for element scatterers + from . import scatter_func_factory + element_scatter_methods = [ + scatter_func_factory.render(e) + for e in elements + ] + element0_interact_path1 = element_scatter_methods[0]['interact_path1'] + element1_interact_path1 = element_scatter_methods[1]['interact_path1'] + element2_interact_path1 = element_scatter_methods[2]['interact_path1'] + element0_calculate_attenuation = element_scatter_methods[0]['calculate_attenuation'] + element1_calculate_attenuation = element_scatter_methods[1]['calculate_attenuation'] + element2_calculate_attenuation = element_scatter_methods[2]['calculate_attenuation'] + del elements, element_scatter_methods + + @cuda.jit(device=True) + def element_interact_path1(threadindex, rng_states, neutron, element_index): + if element_index == 0: + return element0_interact_path1(threadindex, rng_states, neutron) + if element_index == 1: + return element1_interact_path1(threadindex, rng_states, neutron) + if element_index == 2: + return element2_interact_path1(threadindex, rng_states, neutron) + return none + + @cuda.jit(device=True) + def element_calculate_attenuation(neutron, end, element_index): + if element_index == 0: + return element0_calculate_attenuation(neutron, end) + if element_index == 1: + return element1_calculate_attenuation(neutron, end) + if element_index == 2: + return element2_calculate_attenuation(neutron, end) + return 1.0 + + return dict( + element_interact_path1 = element_interact_path1, + element_calculate_attenuation = element_calculate_attenuation, + ) diff --git a/acc/scatterers/composite_scatterer.py b/acc/scatterers/composite_scatterer.py new file mode 100644 index 00000000..b1ed5be2 --- /dev/null +++ b/acc/scatterers/composite_scatterer.py @@ -0,0 +1,296 @@ +# -*- python -*- +# +# Jiao Lin +# + +import os +from math import sqrt, exp +import numpy as np, numba +from numba import cuda, void, int64 +from numba.core import config +from numba.cuda.random import xoroshiro128p_type + +from .. import test +from .._numba import xoroshiro128p_uniform_float32 +from .interaction_types import absorption, scattering, none +from ..neutron import absorb, prop_dt_inplace, clone, is_moving +from ..geometry.locate import inside, outside, onborder + + +from ..config import get_numba_floattype +NB_FLOAT = get_numba_floattype() + +def factory(composite): + elements = composite.elements() + N = len(elements) + mod = _importModule(N) + element_methods = mod.createHelperMethodsForScatter(composite) + element_interact_path1 = element_methods["element_interact_path1"] + element_calculate_attenuation = element_methods["element_calculate_attenuation"] + rt_methods = _createRayTracingMethods(composite) + find_1st_hit = rt_methods['find_1st_hit'] + locate = rt_methods['locate'] + intersect = rt_methods['intersect'] + propagate_out = rt_methods['propagate_out'] + propagate_to_next_exiting_surface = rt_methods['propagate_to_next_exiting_surface'] + is_exiting = rt_methods["is_exiting"] + is_exiting_union = rt_methods["is_exiting_union"] + propagate_to_next_incident_surface_union = rt_methods["propagate_to_next_incident_surface_union"] + + @cuda.jit(device=True) + def _scatter(threadindex, rng_states, neutron, tmp_neutron, tmp_ts): + if not is_moving(neutron): + absorb(neutron) + return + if is_exiting(neutron): + return + if is_exiting_union(neutron): + propagate_to_next_exiting_surface(neutron) + return + while True: + itype = _interact_path1(threadindex, rng_states, neutron, tmp_neutron) + if itype == absorption: + return + if itype == scattering: + x,y,z, vx,vy,vz = neutron[:6] + Nintersect = intersect(x,y,z, vx,vy,vz, tmp_ts) + if Nintersect == 0: + propagate_out(neutron) + return + clone(neutron, tmp_neutron) + propagate_out(neutron) + att = calculate_attenuation(tmp_neutron, neutron[:3]) + neutron[-1] *= att + return + # itype == none + if is_exiting(neutron): + return + if is_exiting_union(neutron): + propagate_to_next_exiting_surface(neutron) + return + propagate_to_next_incident_surface_union(neutron) + return + + @cuda.jit(device=True) + def _interact_path1(threadindex, rng_states, neutron, tmp_neutron): + x,y,z = neutron[:3] + vx,vy,vz = neutron[3:6] + scatterer_index = find_1st_hit(x,y,z, vx,vy,vz) + if scatterer_index < 0 or scatterer_index > N: + propagate_to_next_exiting_surface(neutron) + return none + clone(neutron, tmp_neutron) + # global2local(tmp_neutron, scatterer_index) + interaction = element_interact_path1(threadindex, rng_states, tmp_neutron, scatterer_index) + # local2global(tmp_neutron, scatterer_index) + clone(tmp_neutron, neutron) + if interaction == absorption: + absorb(neutron) + return interaction + # nothing happened + if interaction == none: + # if neutron is not inside, we are done + x,y,z = neutron[:3] + if locate(x,y,z)!=inside: + return interaction + # otherwise, interact again + # return _interact_path1(threadindex, rng_states, neutron, tmp_neutron) + return interaction + # interaction must be scatter + clone(neutron, tmp_neutron) + # propagate to surface if necessary + x,y,z = neutron[:3] + if locate(x,y,z) == inside: + propagate_to_next_exiting_surface(neutron) + # apply attenuation + att = calculate_attenuation(tmp_neutron, neutron[:3]) + neutron[-1] *= att + return scattering + + @cuda.jit(device=True) + def _calculate_attenuation(neutron, end, tmp_neutron): + ret = 1.0 + for i in range(N): + clone(neutron, tmp_neutron) + # global2local(tmp_neutron, i) + ret *= element_calculate_attenuation(neutron, end, i) + return ret + + from ..geometry.arrow_intersect import max_intersections + if test.USE_CUDASIM: + @cuda.jit(device=True, inline=True) + def interact_path1(threadindex, rng_states, neutron): + tmp_neutron = np.zeros(10, dtype=float) + return _interact_path1(threadindex, rng_states, neutron, tmp_neutron) + @cuda.jit(device=True, inline=True) + def calculate_attenuation(neutron, end): + tmp_neutron = np.zeros(10, dtype=float) + return _calculate_attenuation(neutron, end, tmp_neutron) + @cuda.jit(device=True, inline=True) + def scatter(threadindex, rng_states, neutron): + tmp_neutron = np.zeros(10, dtype=float) + tmp_ts = np.zeros(max_intersections, dtype=float) + return _scatter(threadindex, rng_states, neutron, tmp_neutron, tmp_ts) + else: + @cuda.jit(device=True, inline=True) + def interact_path1(threadindex, rng_states, neutron): + tmp_neutron = cuda.local.array(10, dtype=numba.float64) + return _interact_path1(threadindex, rng_states, neutron, tmp_neutron) + @cuda.jit(device=True, inline=True) + def calculate_attenuation(neutron, end): + tmp_neutron = cuda.local.array(10, dtype=numba.float64) + return _calculate_attenuation(neutron, end, tmp_neutron) + @cuda.jit( + void(int64, xoroshiro128p_type[:], NB_FLOAT[:]), + device=True, inline=True + ) + def scatter(threadindex, rng_states, neutron): + tmp_neutron = cuda.local.array(10, dtype=numba.float64) + tmp_ts = cuda.local.array(max_intersections, dtype=numba.float64) + return _scatter(threadindex, rng_states, neutron, tmp_neutron, tmp_ts) + return dict( + scatter = scatter, + interact_path1 = interact_path1, + calculate_attenuation = calculate_attenuation, + ) + + +class _Union: + def identify(self, visitor): return visitor.onUnion(self) + +def _createRayTracingMethods(composite): + """create ray-tracing methods to deal with the overall shape of the composite + """ + # methods regarding the overall shape + shape = composite.shape() + from ..geometry import arrow_intersect + intersect = arrow_intersect.arrow_intersect_func_factory.render(shape) + locate = arrow_intersect.locate_func_factory.render(shape) + del shape + from ..geometry.propagation import makePropagateMethods + propagate_methods = makePropagateMethods(intersect, locate) + ret = propagate_methods + # methods regarding union of element shapes + elements = composite.elements() + shapes = [e.shape() for e in elements] + union = _Union(); union.shapes = shapes + u_intersect = arrow_intersect.arrow_intersect_func_factory.render(union) + u_locate = arrow_intersect.locate_func_factory.render(union) + u_propagate_methods = makePropagateMethods(u_intersect, u_locate) + # find_1st_hit + from mcvine.acc.geometry.composite import get_find_1st_hit + ret.update( + locate = locate, + intersect = intersect, + find_1st_hit = get_find_1st_hit(shapes), + is_exiting_union = u_propagate_methods['is_exiting'], + propagate_to_next_incident_surface_union = u_propagate_methods['propagate_to_next_incident_surface'], + ) + return ret + +def _importModule(N): + mod = _makeModule(N) + import imp + return imp.load_source(f"scatterer_composite_{N}", mod) + +def _makeModule(N, overwrite=False): + "make a python module for cuda device methods for composite scatterer with N elements" + modulepath = coder.getModule("composite_scatterer", N) + if os.path.exists(modulepath) and not overwrite: + return modulepath + indent = 4*' ' + imports = """import os, numpy as np, numba +from numba import cuda +from numba.core import config + +from mcvine.acc.scatterers.interaction_types import none +from mcvine.acc import test +""".splitlines() + createHelperMethodsForScatter = _Coder_createHelperMethodsForScatter(N, indent)() + lines = imports + [''] + createHelperMethodsForScatter + with open(modulepath, 'wt') as ostream: + ostream.write("\n".join(lines)) + return modulepath + +from .._numba import coder +class _Coder_createHelperMethodsForScatter: + + def __init__(self, N, indent=4*' '): + self.N = N + self.indent = indent + return + + def __call__(self): + N, indent = self.N, self.indent + header = f"""# elements +elements = composite.elements() +nelements = len(elements) +# methods for element scatterers +from mcvine.acc.scatterers import scatter_func_factory +element_scatter_methods = [ + scatter_func_factory.render(e) + for e in elements +] +""".splitlines() + interact_path1_loop = coder.unrollLoop( + N = N, + indent = '', + in_loop = ["element{i}_interact_path1 = element_scatter_methods[{i}]['interact_path1']"], + ) + calculate_attenuation_loop = coder.unrollLoop( + N = N, + indent = '', + in_loop = ["element{i}_calculate_attenuation = element_scatter_methods[{i}]['calculate_attenuation']"], + ) + footer = """ +return dict( + element_interact_path1 = element_interact_path1, + element_calculate_attenuation = element_calculate_attenuation, +) +""".splitlines() + body = ( + header + + interact_path1_loop + + calculate_attenuation_loop + + [''] + + self.element_interact_path1() + + self.element_calculate_attenuation() + + footer + ) + add_indent = lambda lines, n: [indent*n+l for l in lines] + return ( + ["def createHelperMethodsForScatter(composite):"] + + add_indent(body, 1) + ) + + def element_interact_path1(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def element_interact_path1(threadindex, rng_states, neutron, element_index):" + ] + loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = [ + "if element_index == {i}:", + indent + "return element{i}_interact_path1(threadindex, rng_states, neutron)", + ] + ) + return header + loop + [indent + 'return none'] + + def element_calculate_attenuation(self): + N, indent = self.N, self.indent + header = [ + "@cuda.jit(device=True)", + "def element_calculate_attenuation(neutron, end, element_index):" + ] + loop = coder.unrollLoop( + N = N, + indent = indent, + in_loop = [ + "if element_index == {i}:", + indent + "return element{i}_calculate_attenuation(neutron, end)", + ] + ) + return header + loop + [indent + 'return 1.0'] diff --git a/acc/scatterers/homogeneous_scatterer.py b/acc/scatterers/homogeneous_scatterer.py new file mode 100644 index 00000000..10f2a8f6 --- /dev/null +++ b/acc/scatterers/homogeneous_scatterer.py @@ -0,0 +1,123 @@ +# -*- python -*- +# +# Jiao Lin +# + +import numpy as np, numba +from numba import cuda, void, int64 +from mcvine.acc._numba import xoroshiro128p_uniform_float32 +from math import sqrt, exp + +from .interaction_types import absorption, scattering, none +from .. import test +from ..vec3 import ( + subtract as v3_subtract, + length as v3_length, + add as v3_add, + add_inplace as v3_add_inplace, + scale as v3_scale, + copy as v3_copy, +) +from ..neutron import absorb, prop_dt_inplace +from ..geometry.arrow_intersect import max_intersections +from ..geometry.locate import inside, outside, onborder +from ..geometry.propagation import makePropagateMethods + + +from numba.core import config +if not config.ENABLE_CUDASIM: + from numba.cuda.compiler import Dispatcher, DeviceFunction + +from ..config import get_numba_floattype +NB_FLOAT = get_numba_floattype() + + +def factory(shape, kernel, mcweights, packing_factor): + w_absorption, w_scattering, w_transmission = mcweights + from ..geometry import arrow_intersect + intersect = arrow_intersect.arrow_intersect_func_factory.render(shape) + locate = arrow_intersect.locate_func_factory.render(shape) + propagate_methods = makePropagateMethods(intersect, locate) + propagate_to_next_incident_surface = propagate_methods['propagate_to_next_incident_surface'] + propagate_to_next_exiting_surface = propagate_methods['propagate_to_next_exiting_surface'] + tof_before_first_exit = propagate_methods['tof_before_first_exit'] + forward_distance_in_shape = propagate_methods["forward_distance_in_shape"] + from ..kernels import scatter_func_factory + scatter, calc_scattering_coeff, absorb, calc_absorption_coeff = scatter_func_factory.render(kernel) + @cuda.jit(device=True) + def calculate_attenuation(neutron, end): + length = forward_distance_in_shape(neutron, end) + sigma = calc_scattering_coeff(neutron) * packing_factor + mu = calc_absorption_coeff(neutron) * packing_factor + return exp( - (mu+sigma) * length ) + + @cuda.jit(device=True) + def _interact_path1(threadindex, rng_states, neutron, tmp_neutron): + x, y, z, vx, vy, vz = neutron[:6] + loc = locate(x,y,z) + if loc != inside: + # propagate to the front surface + propagate_to_next_incident_surface(neutron) + # distance to travel across and leave for the first time + tof = tof_before_first_exit(neutron) + velocity = sqrt(vx*vx+vy*vy+vz*vz) + distance = velocity * tof + # scattering and absorption coeffs + sigma = calc_scattering_coeff(neutron) + mu = calc_absorption_coeff(neutron) + # probability of three interaction types happening + transmission_prob = exp( -(mu+sigma)*distance ) + absorption_prob = (1-transmission_prob)*(mu/(mu+sigma)) + # scattering_prob = (1-transmission_prob)*(sigma/(mu+sigma)) + # + # toss a dice and decide whether we should do transmission, absorption, + # or scattering + transmission_mark = w_transmission + absorption_mark = transmission_mark + w_absorption + sum_of_weights = absorption_mark + w_scattering + r = xoroshiro128p_uniform_float32(rng_states, threadindex) * sum_of_weights + if (r < transmission_mark) : + # transmission + propagate_to_next_exiting_surface( neutron ) + neutron[-1] *= transmission_prob * (sum_of_weights/w_transmission) + return none + if (r >= transmission_mark and r < absorption_mark ): + # absorption + x = xoroshiro128p_uniform_float32(rng_states, threadindex) * distance + prob = mu * distance * exp( -(mu+sigma) * x ) + neutron[-1] *= prob * (sum_of_weights/w_absorption) + prop_dt_inplace( neutron, x/velocity ) + absorb( threadindex, rng_states, neutron ) + neutron[-1] = -1 + return absorption + # scattering + x = xoroshiro128p_uniform_float32(rng_states, threadindex) * distance + atten = exp( -(mu+sigma) * x ) + prob = distance * atten + prob *= sum_of_weights/w_scattering + neutron[-1] *= prob + prop_dt_inplace( neutron, x/velocity ) + scatter( threadindex, rng_states, neutron ) + neutron[-1] *= packing_factor + if neutron[-1] <=0: + return absorption + for i in range(10): + tmp_neutron[i] = neutron[i] + propagate_to_next_exiting_surface( neutron ) + atten2 = calculate_attenuation( tmp_neutron, neutron[:3] ); + neutron[-1] *= atten2; + return scattering + if test.USE_CUDASIM: + @cuda.jit(device=True, inline=True) + def interact_path1(threadindex, rng_states, neutron): + tmp_neutron = np.zeros(10, dtype=float) + _interact_path1(threadindex, rng_states, neutron, tmp_neutron) + else: + @cuda.jit(device=True, inline=True) + def interact_path1(threadindex, rng_states, neutron): + tmp_neutron = cuda.local.array(10, dtype=numba.float64) + _interact_path1(threadindex, rng_states, neutron, tmp_neutron) + return dict( + interact_path1 = interact_path1, + calculate_attenuation = calculate_attenuation, + ) diff --git a/acc/scatterers/interaction_types.py b/acc/scatterers/interaction_types.py new file mode 100644 index 00000000..0ccc30f9 --- /dev/null +++ b/acc/scatterers/interaction_types.py @@ -0,0 +1,3 @@ +absorption = 0 +scattering = 1 +none = 2 diff --git a/acc/vec3.py b/acc/vec3.py index 38d5f16d..ab811b6e 100644 --- a/acc/vec3.py +++ b/acc/vec3.py @@ -27,6 +27,13 @@ def length2(v): """ return v[0]*v[0]+v[1]*v[1]+v[2]*v[2] +@cuda.jit(device=True, inline=True) +def distance(v1, v2): + x = v1[0]-v2[0] + y = v1[1]-v2[1] + z = v1[2]-v2[2] + return math.sqrt(x*x+y*y+z*z) + @cuda.jit(device=True, inline=True) def normalize(v): l = length(v) @@ -47,6 +54,13 @@ def copy(v1, v2): v2[2]=v1[2] return +@cuda.jit(device=True, inline=True) +def add_inplace(v1, v2): # v1+=v2 + v1[0]+=v2[0] + v1[1]+=v2[1] + v1[2]+=v2[2] + return + @cuda.jit(device=True, inline=True) def add(v1, v2, v3): v3[0]=v1[0]+v2[0] diff --git a/setup.py b/setup.py index de29f641..f13e6bd1 100644 --- a/setup.py +++ b/setup.py @@ -24,6 +24,7 @@ 'mcvine.acc.components.optics', 'mcvine.acc.components.monitors', 'mcvine.acc.components.samples', + 'mcvine.acc.scatterers', ], package_dir = { 'mcvine.acc': "acc", @@ -37,6 +38,7 @@ 'mcvine.acc.components.optics' : "acc/components/optics", 'mcvine.acc.components.monitors' : "acc/components/monitors", 'mcvine.acc.components.samples' : "acc/components/samples", + 'mcvine.acc.scatterers' : "acc/scatterers", }, data_files = [], install_requires = [ diff --git a/tests/components/samples/sampleassemblies/sample+2cylinders/Al.xyz b/tests/components/samples/sampleassemblies/sample+2cylinders/Al.xyz new file mode 100644 index 00000000..95e6d938 --- /dev/null +++ b/tests/components/samples/sampleassemblies/sample+2cylinders/Al.xyz @@ -0,0 +1,6 @@ +4 +4.04 0 0 0 4.04 0 0 0 4.04 +Al 0 0 0 +Al 0.5 0.5 0 +Al 0.5 0 0.5 +Al 0 0.5 0.5 diff --git a/tests/components/samples/sampleassemblies/sample+2cylinders/sample-scatterer.xml b/tests/components/samples/sampleassemblies/sample+2cylinders/sample-scatterer.xml new file mode 100644 index 00000000..47b12cc2 --- /dev/null +++ b/tests/components/samples/sampleassemblies/sample+2cylinders/sample-scatterer.xml @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/tests/components/samples/sampleassemblies/sample+2cylinders/sampleassembly.xml b/tests/components/samples/sampleassemblies/sample+2cylinders/sampleassembly.xml new file mode 100644 index 00000000..b4be3c18 --- /dev/null +++ b/tests/components/samples/sampleassemblies/sample+2cylinders/sampleassembly.xml @@ -0,0 +1,47 @@ + + + + + + + + + + + Al + Al.xyz + + + + + + + + + + + + Al + Al.xyz + + + + + + + + + + + + Al + Al.xyz + + + + + + + + + diff --git a/tests/components/samples/sampleassemblies/sample+2cylinders/sampleenv1-scatterer.xml b/tests/components/samples/sampleassemblies/sample+2cylinders/sampleenv1-scatterer.xml new file mode 100644 index 00000000..47b12cc2 --- /dev/null +++ b/tests/components/samples/sampleassemblies/sample+2cylinders/sampleenv1-scatterer.xml @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/tests/components/samples/sampleassemblies/sample+2cylinders/sampleenv2-scatterer.xml b/tests/components/samples/sampleassemblies/sample+2cylinders/sampleenv2-scatterer.xml new file mode 100644 index 00000000..47b12cc2 --- /dev/null +++ b/tests/components/samples/sampleassemblies/sample+2cylinders/sampleenv2-scatterer.xml @@ -0,0 +1,10 @@ + + + + + + + + + + diff --git a/tests/components/samples/test_composite_isotropic_sphere.py b/tests/components/samples/test_composite_isotropic_sphere.py new file mode 100755 index 00000000..d1780996 --- /dev/null +++ b/tests/components/samples/test_composite_isotropic_sphere.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python + +import os, shutil +import pytest +from mcvine.acc import test +from mcvine import run_script + +thisdir = os.path.dirname(__file__) + + +def test_1(): + instr = os.path.join(thisdir, "test_instrument_with_samplefromxml.py") + samplexml = "sampleassemblies/isotropic_sphere/sampleassembly.xml" + outdir = 'out.test_instrument_source_composite_monitor_isotropic_sphere' + if os.path.exists(outdir): shutil.rmtree(outdir) + ncount = 1e5 + run_script.run1( + instr, outdir, + ncount=ncount, + samplexml = samplexml, + factory = "mcvine.acc.components.samples.composite.sampleassembly_from_xml", + ) + return + +@pytest.mark.skipif(not test.USE_CUDA, reason='No CUDA') +def test_compare_mcvine(num_neutrons=int(1e7), debug=False, interactive=False): + """ + Tests the acc cpu implementation of a straight guide against mcvine + """ + instr = os.path.join(thisdir, "test_instrument_with_samplefromxml.py") + samplexml = "sampleassemblies/isotropic_sphere/sampleassembly.xml", + from mcvine.acc.test.compare_acc_nonacc import compare_acc_nonacc + compare_acc_nonacc( + "isotropic_sphere", + ["psd_4pi"], + {"float32": 4e-10, "float64": 4e-10}, + num_neutrons, debug, + instr = instr, + interactive=interactive, + acc_component_spec = dict( + samplexml=samplexml, + factory = "mcvine.acc.components.samples.composite.sampleassembly_from_xml", + ), + nonacc_component_spec = dict( + samplexml=samplexml, + factory="mcvine.components.samples.SampleAssemblyFromXml", + ), + ) + +def psd_monitor_4pi(): + from mcvine.acc.components.monitors.psd_monitor_4pi import PSD_monitor_4Pi + return PSD_monitor_4Pi( + "mon", + nphi=30, ntheta=30, radius=3, + filename = "psd_4pi.h5", + ) + +@pytest.mark.skipif(not test.USE_CUDA, reason='No CUDA') +def test_acc_run_script(ncount = 1e6): + # instr = os.path.join(thisdir, "acc_ss_instrument.py") + instr = os.path.join(thisdir, "acc_composite_isotropic_sphere_instrument.py") + outdir = 'out.composite_isotropic_sphere-acc_run_script' + ncount = int(ncount) + from mcvine.acc import run_script + run_script.run( + instr, outdir, ncount=ncount, + monitor_factory=psd_monitor_4pi, + z_sample = 1., + ) + return + +def main(): + import journal + journal.info("instrument").activate() + # test_1() + # test_compare_mcvine(num_neutrons=int(100), interactive=True, debug=True) + # test_compare_mcvine(num_neutrons=int(1e7), interactive=True) + test_acc_run_script(ncount=1e6) + return + + +if __name__ == '__main__': main() diff --git a/tests/components/samples/test_homogeneous_multiple_scatterer.py b/tests/components/samples/test_homogeneous_multiple_scatterer.py index 2a1c2c39..2c569eeb 100755 --- a/tests/components/samples/test_homogeneous_multiple_scatterer.py +++ b/tests/components/samples/test_homogeneous_multiple_scatterer.py @@ -63,7 +63,7 @@ def run_kernel(rng_states, out_neutrons, N, neutron): rng_states = create_xoroshiro128p_states(1, seed=0) N = np.zeros(1, dtype=int) rt = run_kernel[1,1](rng_states, out_neutrons, N, neutron) - assert N[0] == HMS.max_ms_loops_path1 + 1 + assert N[0] == HMS.max_ms_loops_path1 print(out_neutrons[:N[0]]) return diff --git a/tests/components/samples/test_instrument_with_samplefromxml.py b/tests/components/samples/test_instrument_with_samplefromxml.py new file mode 100644 index 00000000..4d5cd341 --- /dev/null +++ b/tests/components/samples/test_instrument_with_samplefromxml.py @@ -0,0 +1,16 @@ +#!/usr/bin/env python + +import os +thisdir = os.path.dirname(__file__) + +from test_sample_instrument_factory import construct + +def instrument(samplexml, factory="mcvine.components.samples.SampleAssemblyFromXml"): + tokens = factory.split('.') + module = '.'.join(tokens[:-1]) + method = tokens[-1] + import importlib + module = importlib.import_module(module) + factory = getattr(module, method) + target = factory(name="sample", samplexml=samplexml) + return construct(target, size=0.) diff --git a/tests/geometry/test_arrow_intersect.py b/tests/geometry/test_arrow_intersect.py index 696c30a0..23de45e6 100755 --- a/tests/geometry/test_arrow_intersect.py +++ b/tests/geometry/test_arrow_intersect.py @@ -58,6 +58,23 @@ def test_insert_into_sorted_list(): assert N==6 np.testing.assert_array_equal(l[:N], [0., 1., 2., 3., 4., 8.]) + l = np.arange(4.) + d = 8 + N = arrow_intersect.insert_into_sorted_list(d, l, 4) + assert N==4 + np.testing.assert_array_equal(l[:N], [0., 1., 2., 3.]) + + l = np.arange(4.) + d = -1 + N = arrow_intersect.insert_into_sorted_list(d, l, 4) + assert N==4 + np.testing.assert_array_equal(l[:N], [-1., 0., 1., 2.]) + + l = np.arange(4.) + d = 1.5 + N = arrow_intersect.insert_into_sorted_list(d, l, 4) + assert N==4 + np.testing.assert_array_equal(l[:N], [0., 1., 1.5, 2.]) return def main(): diff --git a/tests/geometry/test_composite_3_elements.py b/tests/geometry/test_composite_3_elements.py new file mode 100644 index 00000000..f379aca3 --- /dev/null +++ b/tests/geometry/test_composite_3_elements.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python + +import pytest, os, time +import math, numpy as np +from numba import cuda + +from instrument.nixml import parse_file +from mcvine.acc import test +from mcvine.acc.geometry import locate, location, arrow_intersect +from mcvine.acc.geometry.composite_3 import ( + createRayTracingMethods_NonOverlappingShapes as createRTMethods, + createUnionLocateMethod +) +from mcvine.acc.geometry.composite import _make_find_1st_hit + +thisdir = os.path.dirname(__file__) + +# device functions can be tested with CUDASIM only +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_locate(): + union = parse_file(os.path.join(thisdir, 'union_three_elements.xml'))[0] + shapes = union.shapes + locate_u3 = createUnionLocateMethod(shapes) + assert locate_u3(0, 0, 0) == locate.inside + assert locate_u3(0.025, 0, 0) == locate.onborder + assert locate_u3(0.099, 0, 0) == locate.onborder + assert locate_u3(0.1, 0, 0) == locate.onborder + assert locate_u3(0.199, 0, 0) == locate.onborder + assert locate_u3(0.2, 0, 0) == locate.onborder + assert locate_u3(0.3, 0, 0) == locate.outside + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_intersect_all(): + union = parse_file(os.path.join(thisdir, 'union_three_elements.xml'))[0] + shapes = union.shapes + methods = createRTMethods(shapes) + intersect_all = methods['intersect_all'] + ts = np.zeros(11) + assert intersect_all(0.,0.,0., 1.,0.,0., ts) == 10 + expected = np.array([-0.2,-0.199, -0.1,-0.099, -0.025, 0.025, 0.099,0.1, 0.199,0.2, 0]) + assert np.allclose(ts, expected) + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_forward_intersect_all(): + union = parse_file(os.path.join(thisdir, 'union_three_elements.xml'))[0] + shapes = union.shapes + methods = createRTMethods(shapes) + forward_intersect_all = methods['forward_intersect_all'] + ts = np.zeros(6) + assert forward_intersect_all(0.,0.,0., 1.,0.,0., ts) == 5 + expected = np.array([0.025,0.099,0.1, 0.199,0.2, 0]) + assert np.allclose(ts, expected) + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_find_1st_hit(): + union = parse_file(os.path.join(thisdir, 'union_three_elements.xml'))[0] + shapes = union.shapes + methods = createRTMethods(shapes) + find_1st_hit = _make_find_1st_hit(**methods) + assert find_1st_hit(-0.25,0.,0., 1.,0.,0.) == 0 + assert find_1st_hit(-0.15,0.,0., 1.,0.,0.) == 1 + assert find_1st_hit(-0.05,0.,0., 1.,0.,0.) == 2 + assert find_1st_hit(0.05,0.,0., 1.,0.,0.) == 1 + assert find_1st_hit(0.15,0.,0., 1.,0.,0.) == 0 + assert find_1st_hit(0.25,0.,0., 1.,0.,0.) == -1 diff --git a/tests/geometry/test_composite_shape.py b/tests/geometry/test_composite_shape.py new file mode 100644 index 00000000..558aa396 --- /dev/null +++ b/tests/geometry/test_composite_shape.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python + +import pytest, os, time +import math, numpy as np +from numba import cuda + +from instrument.nixml import parse_file +from mcvine.acc import test +from mcvine.acc.geometry import locate +from mcvine.acc.geometry.composite import _make_find_1st_hit + +thisdir = os.path.dirname(__file__) + +def test_makeModule(): + from mcvine.acc.geometry.composite import _makeModule + _makeModule(4, overwrite=True) + return + +@pytest.fixture +def shapes(): + union = parse_file(os.path.join(thisdir, 'union_three_elements.xml'))[0] + return union.shapes + +@pytest.fixture +def ray_tracing_methods(shapes): + from mcvine.acc.geometry.composite import _importModule + mod = _importModule(len(shapes)) + methods = mod.createRayTracingMethods_NonOverlappingShapes(shapes) + methods['union_locate'] = mod.createUnionLocateMethod(shapes) + return methods + +# device functions can be tested with CUDASIM only +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_locate(ray_tracing_methods): + locate_u3 = ray_tracing_methods['union_locate'] + assert locate_u3(0, 0, 0) == locate.inside + assert locate_u3(0.025, 0, 0) == locate.onborder + assert locate_u3(0.099, 0, 0) == locate.onborder + assert locate_u3(0.1, 0, 0) == locate.onborder + assert locate_u3(0.199, 0, 0) == locate.onborder + assert locate_u3(0.2, 0, 0) == locate.onborder + assert locate_u3(0.3, 0, 0) == locate.outside + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_intersect_all(ray_tracing_methods): + intersect_all = ray_tracing_methods['intersect_all'] + ts = np.zeros(11) + assert intersect_all(0.,0.,0., 1.,0.,0., ts) == 10 + expected = np.array([-0.2,-0.199, -0.1,-0.099, -0.025, 0.025, 0.099,0.1, 0.199,0.2, 0]) + assert np.allclose(ts, expected) + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_forward_intersect_all(ray_tracing_methods): + forward_intersect_all = ray_tracing_methods['forward_intersect_all'] + ts = np.zeros(6) + assert forward_intersect_all(0.,0.,0., 1.,0.,0., ts) == 5 + expected = np.array([0.025,0.099,0.1, 0.199,0.2, 0]) + assert np.allclose(ts, expected) + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_find_1st_hit(ray_tracing_methods): + find_1st_hit = _make_find_1st_hit(**ray_tracing_methods) + assert find_1st_hit(-0.25,0.,0., 1.,0.,0.) == 0 + assert find_1st_hit(-0.15,0.,0., 1.,0.,0.) == 1 + assert find_1st_hit(-0.05,0.,0., 1.,0.,0.) == 2 + assert find_1st_hit(0.05,0.,0., 1.,0.,0.) == 1 + assert find_1st_hit(0.15,0.,0., 1.,0.,0.) == 0 + assert find_1st_hit(0.25,0.,0., 1.,0.,0.) == -1 + +@pytest.mark.skipif(not test.USE_CUDA, reason='no CUDA') +def test_cuda(ray_tracing_methods): + find_1st_hit = _make_find_1st_hit(**ray_tracing_methods) + @cuda.jit + def find_1st_hit_kernel(neutrons, out, n_neutrons_per_thread): + N = len(neutrons) + thread_index = cuda.grid(1) + start_index = thread_index*n_neutrons_per_thread + end_index = min(start_index+n_neutrons_per_thread, N) + for i in range(start_index, end_index): + neutron = neutrons[i] + x,y,z, vx,vy,vz = neutron[:6] + out[i] = find_1st_hit(x,y,z, vx,vy,vz) + def run(neutrons, first_hit): + N = len(neutrons) + threads_per_block = 128 + ntotthreads = int(1e5) + nblocks = math.ceil(ntotthreads / threads_per_block) + actual_nthreads = threads_per_block * nblocks + n_neutrons_per_thread = math.ceil(N / actual_nthreads) + find_1st_hit_kernel[nblocks, threads_per_block](neutrons, first_hit, n_neutrons_per_thread) + N = int(100) + neutrons = np.zeros((N, 10)) + neutrons[:, :3] = np.random.random((N, 3)) * 0.1 + neutrons[:, 3:6] = np.random.random((N, 3)) * 1000 + first_hit = np.zeros(N, dtype=int) + run(neutrons, first_hit) + print(first_hit) + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_public_methods(shapes): + from mcvine.acc.geometry.composite import get_find_1st_hit, get_union_locate + get_find_1st_hit(shapes) + get_union_locate(shapes) + return + diff --git a/tests/geometry/test_propagation.py b/tests/geometry/test_propagation.py index 000582dc..ccc6cc7a 100755 --- a/tests/geometry/test_propagation.py +++ b/tests/geometry/test_propagation.py @@ -134,7 +134,7 @@ def test_propagate_to_next_exiting_surface(): [-0.01,0.,0., 1.,0.,0., 0.,0., 0.015, 1.], [-0.01,0.,0., 1.,0.,0., 0.,0., 0.01, 1.], [-0.01,0.,0., 1.,0.,0., 0.,0., 0.005, 1.], - [-0.01,0.,0., 1.,0.,0., 0.,0., 0.01, 1.], + [-0.01,0.,0., 1.,0.,0., 0.,0., 0.0, 1.], [0.02,0.,0., 1.,0.,0., 0.,0., 0.02, 1.], [0.02,0.,0., 1.,0.,0., 0.,0., 0.01, 1.], [0.02,0.,0., 1.,0.,0., 0.,0., 0.005, 1.], @@ -142,6 +142,29 @@ def test_propagate_to_next_exiting_surface(): [0.025,0.,0., 1.,0.,0., 0.,0., 0., 1.], [1.,0.,0., 1.,0.,0., 0.,0., 0., 1.], ]) + for i, (neutron, out) in enumerate(zip(neutrons, expected)): + propagate_to_next_exiting_surface(neutron) + assert np.allclose(neutron, out), f"{i}: {neutron} != {out}" + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_forward_distance_in_shape(): + forward_distance_in_shape = getPropagateMethods()[ + 'forward_distance_in_shape'] + neutrons = np.array([ + [-1.,0.,0., 1.,0.,0., 0.,0., 0., 1.], + ]) + ends = np.array([ + [0.,0.,0.], + [1.,0.,0.], + ]) + expected = np.array([ + 0.01, + 0.02, + ]) + for neutron, end, expected_out in zip(neutrons, ends, expected): + out = forward_distance_in_shape(neutron, end) + assert np.allclose(out, expected_out), f"{out} != {expected_out}" return def main(): @@ -149,6 +172,7 @@ def main(): test_tof_before_first_exit() test_propagate_to_next_incident_surface() test_propagate_to_next_exiting_surface() + test_forward_distance_in_shape() return if __name__ == '__main__': main() diff --git a/tests/geometry/union_three_elements.xml b/tests/geometry/union_three_elements.xml new file mode 100644 index 00000000..b4112562 --- /dev/null +++ b/tests/geometry/union_three_elements.xml @@ -0,0 +1,19 @@ + + + + + + + + + + + + + + + + + + + diff --git a/tests/scatterers/sampleassemblies b/tests/scatterers/sampleassemblies new file mode 120000 index 00000000..e466c9c6 --- /dev/null +++ b/tests/scatterers/sampleassemblies @@ -0,0 +1 @@ +../components/samples/sampleassemblies \ No newline at end of file diff --git a/tests/scatterers/test_composite_scatterer.py b/tests/scatterers/test_composite_scatterer.py new file mode 100644 index 00000000..dc0a7f7f --- /dev/null +++ b/tests/scatterers/test_composite_scatterer.py @@ -0,0 +1,132 @@ +import os +thisdir = os.path.dirname(__file__) +import math, numpy as np +from numba import cuda +import pytest +from mcvine.acc import test + +@pytest.fixture +def composite_example_methods(): + path = os.path.join(thisdir, "sampleassemblies", 'sample+2cylinders', 'sampleassembly.xml') + from mcvine.acc.components.samples import loadScattererComposite + composite = loadScattererComposite(path) + from mcvine.acc.scatterers.composite_scatterer import factory + return factory(composite) + +@pytest.fixture +def composite_example_interact_path1(composite_example_methods): + return composite_example_methods['interact_path1'] + +@pytest.fixture +def composite_example_scatter(composite_example_methods): + return composite_example_methods['scatter'] + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_cudasim_interact_path1(composite_example_interact_path1): + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + threadindex = 0 + rng_states = None + N = 10 #000 + p = 0. + absorbed = 0 + for i in range(N): + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + composite_example_interact_path1(threadindex, rng_states, neutron) + print(neutron) + p1 = neutron[-1] + if p1>0: + p+=p1 + else: + absorbed += 1 + p/=N + absorbed/=N + print(p, absorbed) + assert p < 0.1 + assert absorbed == 0 + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_cudasim_scatter(composite_example_scatter): + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + threadindex = 0 + rng_states = None + N = 10 #000 + p = 0. + absorbed = 0 + for i in range(N): + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + composite_example_scatter(threadindex, rng_states, neutron) + print(neutron) + p1 = neutron[-1] + if p1>0: + p+=p1 + else: + absorbed += 1 + p/=N + absorbed/=N + print(p, absorbed) + assert p < 0.1 + assert absorbed == 0 + return + +@pytest.mark.skipif(not test.USE_CUDA, reason='no CUDA') +def test_cuda_interact_path1(composite_example_interact_path1): + @cuda.jit + def interact_path1_kernel(rng_states, neutrons, n_neutrons_per_thread): + N = len(neutrons) + thread_index = cuda.grid(1) + start_index = thread_index*n_neutrons_per_thread + end_index = min(start_index+n_neutrons_per_thread, N) + for i in range(start_index, end_index): + composite_example_interact_path1(thread_index, rng_states, neutrons[i]) + def run(neutrons): + N = len(neutrons) + threads_per_block = 128 + ntotthreads = int(1e5) + nblocks = math.ceil(ntotthreads / threads_per_block) + actual_nthreads = threads_per_block * nblocks + n_neutrons_per_thread = math.ceil(N / actual_nthreads) + from numba.cuda.random import create_xoroshiro128p_states + from mcvine.acc.config import rng_seed + rng_states = create_xoroshiro128p_states(actual_nthreads, seed=rng_seed) + interact_path1_kernel[nblocks, threads_per_block](rng_states, neutrons, n_neutrons_per_thread) + N = int(10) + neutrons = np.zeros((N, 10)) + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + neutrons[:] = neutron + run(neutrons) + print(neutrons) + return + +@pytest.mark.skipif(not test.USE_CUDA, reason='no CUDA') +def test_cuda_scatter(composite_example_scatter): + @cuda.jit + def scatter_kernel(rng_states, neutrons, n_neutrons_per_thread): + N = len(neutrons) + thread_index = cuda.grid(1) + start_index = thread_index*n_neutrons_per_thread + end_index = min(start_index+n_neutrons_per_thread, N) + for i in range(start_index, end_index): + composite_example_scatter(thread_index, rng_states, neutrons[i]) + def run(neutrons): + N = len(neutrons) + threads_per_block = 128 + ntotthreads = int(1e5) + nblocks = math.ceil(ntotthreads / threads_per_block) + actual_nthreads = threads_per_block * nblocks + n_neutrons_per_thread = math.ceil(N / actual_nthreads) + from numba.cuda.random import create_xoroshiro128p_states + from mcvine.acc.config import rng_seed + rng_states = create_xoroshiro128p_states(actual_nthreads, seed=rng_seed) + scatter_kernel[nblocks, threads_per_block](rng_states, neutrons, n_neutrons_per_thread) + N = int(10) + neutrons = np.zeros((N, 10)) + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + neutrons[:] = neutron + run(neutrons) + print(neutrons) + return + +def test_makeModule(): + from mcvine.acc.scatterers.composite_scatterer import _makeModule + _makeModule(4, overwrite=True) diff --git a/tests/scatterers/test_homogeneous_scatterer.py b/tests/scatterers/test_homogeneous_scatterer.py new file mode 100644 index 00000000..39fa5f33 --- /dev/null +++ b/tests/scatterers/test_homogeneous_scatterer.py @@ -0,0 +1,78 @@ +import os +thisdir = os.path.dirname(__file__) +import math, numpy as np +from numba import cuda +import pytest +from mcvine.acc import test + +@pytest.fixture +def hs_example(): + path = os.path.join(thisdir, "sampleassemblies", 'isotropic_sphere', 'sampleassembly.xml') + from mcvine.acc.components.samples import loadFirstHomogeneousScatterer + return loadFirstHomogeneousScatterer(path) + +@pytest.fixture +def hs_example_interact_path1(hs_example): + shape = hs_example.shape() + kernel = hs_example.kernel() + mcweights = 1., 1., 1. + packing_factor = 0.6 + from mcvine.acc.scatterers.homogeneous_scatterer import factory + methods = factory(shape, kernel, mcweights, packing_factor) + return methods['interact_path1'] + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_homogeneous_scatterer(hs_example_interact_path1): + threadindex = 0 + rng_states = None + N = 10000 + p = 0. + absorbed = 0 + for i in range(N): + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + hs_example_interact_path1(threadindex, rng_states, neutron) + p1 = neutron[-1] + if p1>0: + p+=p1 + else: + absorbed += 1 + p/=N + absorbed/=N + assert p<.8 and p>.7 + assert absorbed < 0.36 and absorbed > 0.3 + return + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test_scatter_func_factory(hs_example): + from mcvine.acc.scatterers import scatter_func_factory + methods = scatter_func_factory.render(hs_example) + return + +@pytest.mark.skipif(not test.USE_CUDA, reason='no CUDA') +def test_homogeneous_scatterer_cuda(hs_example_interact_path1): + @cuda.jit + def interact_path1_kernel(rng_states, neutrons, n_neutrons_per_thread): + N = len(neutrons) + thread_index = cuda.grid(1) + start_index = thread_index*n_neutrons_per_thread + end_index = min(start_index+n_neutrons_per_thread, N) + for i in range(start_index, end_index): + hs_example_interact_path1(thread_index, rng_states, neutrons[i]) + def run(neutrons): + N = len(neutrons) + threads_per_block = 512 + ntotthreads = int(1e5) + nblocks = math.ceil(ntotthreads / threads_per_block) + actual_nthreads = threads_per_block * nblocks + n_neutrons_per_thread = math.ceil(N / actual_nthreads) + from numba.cuda.random import create_xoroshiro128p_states + from mcvine.acc.config import rng_seed + rng_states = create_xoroshiro128p_states(actual_nthreads, seed=rng_seed) + interact_path1_kernel[nblocks, threads_per_block](rng_states, neutrons, n_neutrons_per_thread) + N = int(1e5) + neutrons = np.zeros((N, 10)) + neutron = np.array([0.0,0,0, 0,0,1000, 0,0, 0, 1.]) + neutrons[:] = neutron + run(neutrons) + return + diff --git a/tests/scatterers/test_scatter_func_factory.py b/tests/scatterers/test_scatter_func_factory.py new file mode 100644 index 00000000..dee61bd9 --- /dev/null +++ b/tests/scatterers/test_scatter_func_factory.py @@ -0,0 +1,14 @@ +import os +thisdir = os.path.dirname(__file__) +import numpy as np +import pytest +from mcvine.acc import test + +@pytest.mark.skipif(not test.USE_CUDASIM, reason='no CUDASIM') +def test(): + path = os.path.join(thisdir, "sampleassemblies", 'sample+2cylinders', 'sampleassembly.xml') + from mcvine.acc.components.samples import loadScattererComposite + composite = loadScattererComposite(path) + from mcvine.acc.scatterers import scatter_func_factory + scatter_func_factory.render(composite) + return