Support STS CHESS instrument

acc/components/__init__.py

@@ -0,0 +1 @@
+from . import monitors, optics, samples, sources

acc/components/monitors/__init__.py

@@ -0,0 +1,5 @@
+from .iqe_monitor import IQE_monitor
+from .posdiv_monitor import PosDiv_monitor
+from .psd_monitor import PSD_monitor
+from .psd_monitor_4pi import PSD_monitor_4Pi
+from .wavelength_monitor import Wavelength_monitor

acc/components/optics/__init__.py

@@ -0,0 +1,7 @@
+from .arm import Arm
+from .beamstop import Beamstop
+from .guide import Guide
+# from .guide_anyshape_gravity import Guide_anyshape_gravity
+from .guide_anyshape import Guide_anyshape
+from .guide_tapering import Guide as Guide_tapering
+from .slit import Slit

acc/components/optics/geometry2d.py

@@ -11,7 +11,7 @@
 # modified from
 # https://stackoverflow.com/questions/1119627/how-to-test-if-a-point-is-inside-of-a-convex-polygon-in-2d-integer-coordinates#:~:text=If%20it%20is%20convex%2C%20a,traversed%20in%20the%20same%20order)
 @cuda.jit(device=True)
-def inside_convex_polygon(point, vertices):
+def inside_convex_polygon(point, vertices, l_epsilon):
     # previous_side = get_side(v_sub(vertices[1], vertices[0]), v_sub(point, vertices[0]))
     previous_sign = cosine_sign(v_sub(vertices[1], vertices[0]), v_sub(point, vertices[0]))
     previous_side = previous_sign <= 0
@@ -21,11 +21,11 @@ def inside_convex_polygon(point, vertices):
         affine_segment = v_sub(b, a)
         affine_point = v_sub(point, a)
         sign = cosine_sign(affine_segment, affine_point)
-        if fabs(sign) < 1E-14:
+        if fabs(sign) < l_epsilon:
             continue
         # current_side = get_side(affine_segment, affine_point)
         current_side = sign <= 0
-        if fabs(previous_sign) < 1E-14:
+        if fabs(previous_sign) < l_epsilon:
             previous_side = current_side
             continue
         if previous_side != current_side:

acc/components/optics/guide_anyshape.py

@@ -1,6 +1,8 @@
 # -*- python -*-
 #
 
+import logging
+logger = logging.getLogger(__name__)
 from math import sqrt, fabs
 import numpy as np, numba
 from numba import cuda, void
@@ -11,14 +13,16 @@
 from ...config import get_numba_floattype
 NB_FLOAT = get_numba_floattype()
 from ._guide_utils import calc_reflectivity
-from ...neutron import absorb
+from ...neutron import absorb, clone, prop_dt_inplace
 from ...geometry._utils import insert_into_sorted_list_with_indexes
 from .geometry2d import inside_convex_polygon
 from ... import vec3
 
-max_bounces = 1000
-max_numfaces = 100
-
+max_bounces = 30
+max_numfaces = 5000
+history_size_for_bounces_in_epsilon_neighborhood = 10 # This normally would not exceed 3 for typical guide design.
+t_epsilon = 1E-14
+l_epsilon = 1E-11
 
 @cuda.jit(NB_FLOAT(NB_FLOAT[:], NB_FLOAT[:], NB_FLOAT[:], NB_FLOAT[:, :],  NB_FLOAT[:]),
           device=True, inline=True)
@@ -85,60 +89,134 @@ def load_scaled_centered_faces(path, xwidth=0, yheight=0, zdepth=0, center=False
     faces = np.array([[vertices[i] for i in face] for face in faces])
     return faces
 
+@cuda.jit(numba.boolean(NB_FLOAT[:], NB_FLOAT[:],  NB_FLOAT[:, :], NB_FLOAT[:, :]),
+          device=True, inline=True)
+def likely_inside_face(postmp, face_center, face_uvecs, face2d_bounds):
+    vec3.subtract(postmp, face_center, postmp)
+    ex = face_uvecs[0] 
+    ey = face_uvecs[1] 
+    x = vec3.dot(ex, postmp)
+    y = vec3.dot(ey, postmp)
+    return (
+        (x>face2d_bounds[0, 0]-l_epsilon) and (x<face2d_bounds[1, 0]+l_epsilon)
+        and (y>face2d_bounds[0, 1]-l_epsilon) and (y<face2d_bounds[1, 1]+l_epsilon)
+    )  
+
+@cuda.jit(numba.boolean(NB_FLOAT[:], NB_FLOAT[:],  NB_FLOAT[:]),
+          device=True, inline=True)
+def on_the_other_side(pos, face_center, face_norm):
+    l0 = vec3.dot(face_center, face_norm)
+    l = vec3.dot(pos, face_norm)
+    # logger.debug(f"  on_the_other_side: pos={pos}, face_center={face_center}, face_norm={face_norm}, l={l}, l0={l0}")
+    if l*l0 < 0: return False
+    if abs(l0) < abs(l):
+        # import pdb; pdb.set_trace()
+        return True
+    return False
 
 @cuda.jit(
     void(
         NB_FLOAT[:],
-        NB_FLOAT[:, :, :], NB_FLOAT[:, :], NB_FLOAT[:, :, :], NB_FLOAT[:, :, :],
+        NB_FLOAT[:, :, :], NB_FLOAT[:, :], NB_FLOAT[:, :, :], NB_FLOAT[:, :, :], NB_FLOAT[:, :, :],
         NB_FLOAT, NB_FLOAT, NB_FLOAT, NB_FLOAT, NB_FLOAT,
+        NB_FLOAT[:], numba.int32[:],
         NB_FLOAT[:], NB_FLOAT[:], numba.int32[:],
     ), device=True
 )
 def _propagate(
         in_neutron,
-        faces, centers, unitvecs, faces2d,
+        faces, centers, unitvecs, faces2d, bounds2d,
         R0, Qc, alpha, m, W,
-        tmp1, intersections, face_indexes,
+        intersections, face_indexes,
+        tmpv3, tmp_neutron, tmp_face_hist,
 ):
     nfaces = len(faces)
+    # logger.debug(f"in_neutron={in_neutron}, faces {faces}")
+    N_tmp_face_hist = len(tmp_face_hist)
+    face_hist_start_ind = 0
+    face_hist_size = 0
     for nb in range(max_bounces):
+        # logger.debug(f"bounce {nb}: in_neutron={in_neutron}, faces {faces}")
         # calc intersection with each face and save positive ones in a list with increasing order
         ninter = 0
         for iface in range(nfaces):
-            intersection = intersect_plane(
-                in_neutron[:3], in_neutron[3:6],
-                centers[iface], unitvecs[iface],
-                tmp1
-            )
-            if intersection>0:
-                ninter = insert_into_sorted_list_with_indexes(iface, intersection, face_indexes, intersections, ninter) 
+            found_in_history = False
+            # logger.debug(f"  face {iface}, {face_hist_start_ind}, {face_hist_start_ind+face_hist_size}")
+            for ifh in range(face_hist_start_ind, face_hist_start_ind+face_hist_size):
+                # logger.debug(f"  face {iface}, {ifh}, {tmp_face_hist[ifh]}")
+                if iface == tmp_face_hist[ifh]:
+                    found_in_history = True
+                    break
+            # logger.debug(f"  face {iface} of {nfaces} faces: found_in_history={found_in_history}")
+            if found_in_history: continue
+            x0 = in_neutron[:3]; v0 = in_neutron[3:6]
+            face_center = centers[iface]
+            face_uvecs = unitvecs[iface]
+            intersection = intersect_plane( x0, v0, face_center, face_uvecs, tmpv3)
+            # logger.debug(f"  face {iface}: x0={x0}, v0={v0}, face_center={face_center}, face_uvecs={face_uvecs}, intersection={intersection}")
+            # if intersection<=-t_epsilon:
+            # if intersection<0:
+            if intersection<=-t_epsilon/5:
+                continue
+            face2d_bounds = bounds2d[iface]
+            clone(in_neutron, tmp_neutron)
+            prop_dt_inplace(tmp_neutron, intersection)
+            vec3.copy(tmp_neutron[:3], tmpv3)
+            if not likely_inside_face(tmpv3, face_center, face_uvecs, face2d_bounds):
+                continue
+            # logger.debug(f"  face {iface}: intersection likely inside face")
+            # vec3.copy(tmp_neutron[:3], tmpv3)
+            # if intersection < t_epsilon and on_the_other_side(tmpv3, face_center, face_uvecs[2]):
+            #     continue
+            # logger.debug(f"  face {iface}: intersection on the right side")
+            ninter = insert_into_sorted_list_with_indexes(iface, intersection, face_indexes, intersections, ninter) 
+            # logger.debug(f"  face {iface}: new intersection list {intersections[:ninter]} for faces {face_indexes[:ninter]}")
+        # logger.debug(f"bounce {nb}: ninter={ninter}")
+        # import pdb; pdb.set_trace()
         if not ninter:
             break
         # find the smallest intersection that is inside the mirror 
-        found = False
+        found = -1
         for iinter in range(ninter):
             face_index = face_indexes[iinter]
             intersection = intersections[iinter]
             # calc position of intersection
-            vec3.copy(in_neutron[3:6], tmp1)
-            vec3.scale(tmp1, intersection)
-            vec3.add(tmp1, in_neutron[:3], tmp1)
+            vec3.copy(in_neutron[3:6], tmpv3)
+            vec3.scale(tmpv3, intersection)
+            vec3.add(tmpv3, in_neutron[:3], tmpv3)
             # calc 2d coordinates and use it to check if it is inside the mirror
-            vec3.subtract(tmp1, centers[face_index], tmp1)
+            vec3.subtract(tmpv3, centers[face_index], tmpv3)
             e0 = unitvecs[face_index, 0, :]
             e1 = unitvecs[face_index, 1, :]
             e2 = unitvecs[face_index, 2, :]
             face2d = faces2d[face_index]
-            if inside_convex_polygon((vec3.dot(tmp1, e0), vec3.dot(tmp1, e1)), face2d):
-                found = True
+            # logger.debug(f"check intersection {iinter}: face={face_center}, {face_uvecs[2]}, intersection={intersection}")
+            if inside_convex_polygon((vec3.dot(tmpv3, e0), vec3.dot(tmpv3, e1)), face2d, l_epsilon):
+                found = face_index
+                # logger.debug(f"check intersection {iinter}: found={found}")
+                to_travel_after_bounce = l_epsilon/vec3.length(in_neutron[3:6]) # move a tiny distance
+                if iinter<ninter-1:
+                    next_intersection = intersections[iinter+1]
+                    to_travel_after_bounce = min((next_intersection-intersection)/2, to_travel_after_bounce)
                 break
-        if not found:
+        if found < 0:
             break
+        if intersection<t_epsilon:
+            if face_hist_size < N_tmp_face_hist:
+                tmp_face_hist[(face_hist_start_ind+face_hist_size) % N_tmp_face_hist] = found
+                face_hist_size += 1
+            else:
+                face_hist_start_ind = (face_hist_start_ind+1) % N_tmp_face_hist
+                tmp_face_hist[(face_hist_start_ind-1) % N_tmp_face_hist] = found
+        else:
+            tmp_face_hist[face_hist_start_ind] = found
+            face_hist_size = 1
+
         x, y, z, vx, vy, vz = in_neutron[:6]
         t = in_neutron[-2]
         prob = in_neutron[-1]
         # propagate to intersection
-        intersection -= intersection * 1E-14
+        # intersection -= intersection * 1E-14
         x += vx * intersection
         y += vy * intersection
         z += vz * intersection
@@ -150,14 +228,16 @@ def _propagate(
         if prob <= 0:
             absorb(in_neutron)
             break
-        # tmp1 = velocity change vector
-        vec3.copy(e2, tmp1)
-        vec3.scale(tmp1, vq)
+        # tmpv3 = velocity change vector
+        vec3.copy(e2, tmpv3)
+        vec3.scale(tmpv3, vq)
         # change direction
-        vec3.add(in_neutron[3:6], tmp1, in_neutron[3:6])
+        vec3.add(in_neutron[3:6], tmpv3, in_neutron[3:6])
         in_neutron[:3] = x, y, z
         in_neutron[-2] = t
         in_neutron[-1] = prob
+        prop_dt_inplace(in_neutron, to_travel_after_bounce)
+        # logger.debug(f"bounce {nb}: out_neutron={in_neutron}")
     return
 
 def get_faces_data(geometry, xwidth, yheight, zdepth, center):
@@ -172,8 +252,11 @@ def get_faces_data(geometry, xwidth, yheight, zdepth, center):
         ]
         for face, center, (ex,ey,ez) in zip(faces, centers, unitvecs)
     ]) # nfaces, nverticesperface, 2
-    return faces, centers, unitvecs, faces2d,
-
+    bounds2d = np.array([
+        [np.min(face2d, axis=0), np.max(face2d, axis=0)]
+        for face2d in faces2d
+    ]) # nfaces, 2, 2. bounds2d[iface][0]: minx, miny; bounds2[iface][1]: maxx, maxy
+    return faces, centers, unitvecs, faces2d, bounds2d
 
 from ..ComponentBase import ComponentBase
 class Guide_anyshape(ComponentBase):
@@ -205,16 +288,10 @@ def __init__(
         geometry (str): path of the OFF/PLY geometry file for the guide shape
         """
         self.name = name
-        faces, centers, unitvecs, faces2d = get_faces_data(geometry, xwidth, yheight, zdepth, center)
-        faces2d = np.array([
-            [
-                [np.dot(vertex-center, ex), np.dot(vertex-center, ey)]
-                for vertex in face
-            ]
-            for face, center, (ex,ey,ez) in zip(faces, centers, unitvecs)
-        ]) # nfaces, nverticesperface, 2
+        faces, centers, unitvecs, faces2d, bounds2d = get_faces_data(geometry, xwidth, yheight, zdepth, center)
+        assert len(faces) < max_numfaces
         self.propagate_params = (
-            faces, centers, unitvecs, faces2d,
+            faces, centers, unitvecs, faces2d, bounds2d,
             float(R0), float(Qc), float(alpha), float(m), float(W),
         )
 
@@ -228,22 +305,23 @@ def __init__(
     @cuda.jit(
         void(
             NB_FLOAT[:],
-            NB_FLOAT[:, :, :], NB_FLOAT[:, :], NB_FLOAT[:, :, :], NB_FLOAT[:, :, :],
+            NB_FLOAT[:, :, :], NB_FLOAT[:, :], NB_FLOAT[:, :, :], NB_FLOAT[:, :, :], NB_FLOAT[:, :, :],
             NB_FLOAT, NB_FLOAT, NB_FLOAT, NB_FLOAT, NB_FLOAT,
         ), device=True, inline=True,
     )
     def propagate(
             in_neutron,
-            faces, centers, unitvecs, faces2d,
+            faces, centers, unitvecs, faces2d, bounds2d,
             R0, Qc, alpha, m, W,
     ):
-        tmp1 = cuda.local.array(3, dtype=numba.float64)
-        nfaces = len(faces)
-        assert nfaces < max_numfaces
         intersections = cuda.local.array(max_numfaces, dtype=numba.float64)
         face_indexes = cuda.local.array(max_numfaces, dtype=numba.int32)
+        tmpv3 = cuda.local.array(3, dtype=numba.float64)
+        tmp_neutron = cuda.local.array(10, dtype=numba.float64)
+        tmp_face_hist = cuda.local.array(history_size_for_bounces_in_epsilon_neighborhood, dtype=numba.int32)
         return _propagate(
-            in_neutron, faces, centers, unitvecs, faces2d,
+            in_neutron, faces, centers, unitvecs, faces2d, bounds2d,
             R0, Qc, alpha, m, W,
-            tmp1, intersections, face_indexes,
+            intersections, face_indexes,
+            tmpv3,  tmp_neutron, tmp_face_hist,
         )