Geometric beam transport from intermediate source to sample and detector

• calculate from source to imaging plane with geometric ray instead of assuming a collimated beam impinging on a combined KBS and BOZ lens
• source size effect and KBS included
• display spot size in a table
• provide a setup.py for module installation
• provide documentation

Closes #1 (closed) #3 (closed) #4 (closed)

GeoBeams.py

+# -*- coding: utf-8 -*-
+""" Geometric beam calculator for SCS.
+
+    Copyright (2021) SCS Team.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from sympy.physics.optics import GeometricRay, FreeSpace, ThinLens
+from sympy.solvers import solve
+from sympy import symbols, Matrix, simplify
+from sympy.utilities.lambdify import lambdify
+
+class GeoBeams:
+    """The Beams object to compute beam propagation from source point to sample.
+
+    Inputs
+    ------
+    elems : dict
+        dictionnaries of numerical values for the different parameters.
+    """
+    def __init__(self):
+        #all distances in meters
+        self.elems = {
+            'dIHF': -56, #Intermediate Horizontal Focus position
+            'dEX': -30, #Exit Slit position
+            'dHFM': -3.35, #Horizontal Focusing Mirror position
+            'fHFM': 5.74, #HFM focal length
+            'dVFM': -2, #Vertical Focusing Mirror position
+            'fVFM': 5.05, #VFM focal length
+            'dBOZ': -0.23, #Beam splitting Off axis Zone plate position
+            'fBOZ': 0.25, #BOZ focal length
+            'theta_grating': 0, #grating deflection angle in rad
+            'dSAMZ': 0.0, #Sample position
+            'WH': 0.8e-3, #BOZ horizontal width
+            'WV': 0.8e-3, #BOZ vertical width
+            'offaxis': -0.55e-3, #BOZ center to optical axis
+            'EXw': 100e-6, #Exit Slit width
+            'IHFw': 200e-6, #IHF source width
+            'x1': 0, #horizontal beam height at source
+            'theta1_x': 0, #horizontal beam angle at source
+            'x2': 0, #horizontal beam height at BOZ
+            'y1': 0, #vertical beam height at source
+            'theta1_y': 0, #vertical beam angle at source
+            'y2': 0 #vertical beam height at BOZ
+        }
+        self.compute_eqs()
+
+    def compute_eqs(self):
+        """Compute the beam propagation equation between source and zone plate.
+
+        eq_x, eq_y: tuple, the horizontal x and vertical y equations to
+            compute the initial angle from the source to pass by the zone plate
+            position x2, given the source position at x1.
+        """
+        dEX, dIHF, dVFM, dHFM, dBOZ = symbols('dEX, dIHF, dVFM, dHFM, dBOZ')
+        fVFM, fHFM, fBOZ = symbols('fVFM, fHFM, fBOZ')
+        x1, theta1_x, x2 = symbols('x1, theta1_x, x2')
+        y1, theta1_y, y2 = symbols('y1, theta1_y, y2')
+
+        # horizontal and vertical beam propagation from the
+        # source to the zone plate
+        self.HFM = simplify(FreeSpace(-dHFM+dBOZ)*ThinLens(fHFM)*
+                       FreeSpace(-dIHF+dHFM)*GeometricRay(x1, theta1_x))
+        self.VFM = simplify(FreeSpace(-dVFM+dBOZ)*ThinLens(fVFM)*
+                       FreeSpace(-dEX+dVFM)*GeometricRay(y1, theta1_y))
+        self.HFM_f = lambdify(self.elems.keys(), self.HFM, ['numpy'])
+        self.VFM_f = lambdify(self.elems.keys(), self.VFM, ['numpy'])
+
+
+        # beams for the zone plate first order (0th order is unfocused)
+        self.HBOZ = simplify(ThinLens(fBOZ)*self.HFM)
+        self.VBOZ = simplify(ThinLens(fBOZ)*self.VFM)
+        self.HBOZ_f = lambdify(self.elems.keys(), self.HBOZ, ['numpy'])
+        self.VBOZ_f = lambdify(self.elems.keys(), self.VBOZ, ['numpy'])
+
+        # solve which beam angle at the source theta1 ends on the
+        # zone plate at x2, given x1
+        self.theta1_x = solve(self.HFM[0] - x2, theta1_x)[0]
+        self.theta1_y = solve(self.VFM[0] - y2, theta1_y)[0]
+        self.theta1_x_f = lambdify(self.elems.keys(), self.theta1_x, ['numpy'])
+        self.theta1_y_f = lambdify(self.elems.keys(), self.theta1_y, ['numpy'])
+
+    def path(self):
+        """Compute the horizontal and vertical beam propagation.
+
+        """
+
+        x = []
+        z_x = []
+
+        #source
+        x1 = self.elems['x1']
+        angle1 = self.elems['theta1_x']
+        x += [x1]
+        z_x += [self.elems['dIHF']]
+        ray = GeometricRay(x1, angle1)
+
+        #KBS
+        ray = ThinLens(self.elems['fHFM'])*FreeSpace(-self.elems['dIHF']+self.elems['dHFM'])*ray
+        z_x += [self.elems['dHFM']]
+        x += [ray[0].evalf()]
+
+        #BOZ
+        ray = ThinLens(self.elems['fBOZ'])*FreeSpace(-self.elems['dHFM']+self.elems['dBOZ'])*ray
+        z_x += [self.elems['dBOZ']]
+        x += [ray[0].evalf()]
+
+        #sample
+        ray = FreeSpace(-self.elems['dBOZ']+self.elems['dSAMZ'])*ray
+        z_x += [self.elems['dSAMZ']]
+        x += [ray[0].evalf()]
+
+        y = []
+        z_y = []
+
+        #source
+        y1 = self.elems['y1']
+        angle1 = self.elems['theta1_y']
+        y += [y1]
+        z_y += [self.elems['dEX']]
+        ray = GeometricRay(y1, angle1)
+
+        #KBS
+        ray = ThinLens(self.elems['fVFM'])*FreeSpace(-self.elems['dEX']+self.elems['dVFM'])*ray
+        z_y += [self.elems['dVFM']]
+        y += [ray[0].evalf()]
+
+        #BOZ
+        ray = ThinLens(self.elems['fBOZ'])*FreeSpace(-self.elems['dVFM']+self.elems['dBOZ'])*ray
+        z_y += [self.elems['dBOZ']]
+        y += [ray[0].evalf()]
+
+        #sample
+        ray = FreeSpace(-self.elems['dBOZ']+self.elems['dSAMZ'])*ray
+        z_y += [self.elems['dSAMZ']]
+        y += [ray[0].evalf()]
+
+        return z_x, x, z_y, y
+
+    def plot(self):
+        fig, ax = plt.subplots(2,1, figsize=(6,4), sharex=True, sharey=True)
+
+
+        xmin, xmax = [np.Inf, -np.Inf]
+        ymin, ymax = [np.Inf, -np.Inf]
+        for k,(x2, y2) in enumerate(zip([self.elems['WH']/2, -self.elems['WH']/2],
+                                        np.array([self.elems['WV']/2, -self.elems['WV']/2])
+                                            + self.elems['offaxis'])):
+            for l,(x1, y1) in enumerate(zip([self.elems['IHFw']/2, -self.elems['IHFw']/2],
+                                            [self.elems['EXw']/2, -self.elems['EXw']/2])):
+                c = f'C{3*k+l}'
+
+                self.elems['x1'] = x1
+                self.elems['x2'] = x2
+                self.elems['y1'] = y1
+                self.elems['y2'] = y2
+
+                self.elems['theta1_x'] = self.theta1_x_f(*list(self.elems.values()))
+                self.elems['theta1_y'] = self.theta1_y_f(*list(self.elems.values()))
+
+                z_x, x, z_y, y = self.path()
+                ax[0].plot(z_x, 1e3*np.array(x), c=c)
+                ax[1].plot(z_y, 1e3*np.array(y), c=c)
+
+                if xmin > x[-1]:
+                    xmin = x[-1]
+                if xmax < x[-1]:
+                    xmax = x[-1]
+
+                if ymin > y[-1]:
+                    ymin = y[-1]
+                if ymax < y[-1]:
+                    ymax = y[-1]
+
+        ax[0].set_ylabel('x (mm)')
+        ax[1].set_ylabel('y (mm)')
+        ax[1].set_xlabel('z (m)')
+
+        return fig
+
+    def plane_image(self, p0, n, ZPorder, Gorder=0):
+        """Compute the BOZ image on a plane.
+
+        Inputs
+        ------
+            p0: [x, y, z] point on the imaging plane
+            n: [X, Y, Z] plane normal vector
+            ZPorder: int, order of the zone plate, i.e. 0 or 1
+            Gorder: int, grating order, in [-1, 0, 1]
+        Returns
+        -------
+            list of 4 corners [x,y]
+        """
+        xmin, xmax = [np.Inf, -np.Inf]
+        ymin, ymax = [np.Inf, -np.Inf]
+        for k,(x2, y2) in enumerate(zip(
+                [self.elems['WH']/2, -self.elems['WH']/2],
+                np.array([self.elems['WV']/2, -self.elems['WV']/2])
+                                            + self.elems['offaxis'])):
+            for l,(x1, y1) in enumerate(zip(
+                    [self.elems['IHFw']/2, -self.elems['IHFw']/2],
+                    [self.elems['EXw']/2, -self.elems['EXw']/2])):
+                c = f'C{3*k+l}'
+
+                self.elems['x1'] = x1
+                self.elems['y1'] = y1
+                self.elems['x2'] = x2
+                self.elems['y2'] = y2
+
+                self.elems['theta1_x'] = self.theta1_x_f(
+                    *list(self.elems.values()))
+                self.elems['theta1_y'] = self.theta1_y_f(
+                    *list(self.elems.values()))
+
+                # evaluate HFM/VFM or HBOZ/VBOZ
+                if ZPorder == 1:
+                    bx = self.HBOZ_f(*list(self.elems.values()))
+                    by = self.VBOZ_f(*list(self.elems.values()))
+                elif ZPorder == 0:
+                    bx = self.HFM_f(*list(self.elems.values()))
+                    by = self.VFM_f(*list(self.elems.values()))
+                else:
+                    raise ValueError('ZPorder other than 0 or 1 not implemented')
+
+                # intersection beam with plane
+                l1 = np.array([x2, y2, self.elems['dBOZ']])
+                l12 = np.array([bx[1][0] + Gorder*self.elems['theta_grating'],
+                                by[1][0], 1])
+
+                p = self.LinePlaneIntersection(p0, n, l1, l12=l12)
+
+                if xmin > p[0]:
+                    xmin = p[0]
+                if xmax < p[0]:
+                    xmax = p[0]
+
+                if ymin > p[1]:
+                    ymin = p[1]
+                if ymax < p[1]:
+                    ymax = p[1]
+
+        return np.array([[xmax, ymax], [xmax, ymin],
+                         [xmin, ymin], [xmin, ymax]])
+
+    def LinePlaneIntersection(self, p0, n, l1, *, l2=None, l12=None):
+        """ Calculate the intersection point in space between a line (beam)
+            passing through 2 points l1 and l2 and a (sample) plane passing by
+            p0 with normal n
+
+            l1: [x,y,z] point on line
+            l2: [x,y,z] point on line
+            p0: [x,y,z] point on plane
+            n: plane normal vector
+        """
+        assert not((l2 is None) and (l12 is None)), "Either l2 or l12 must be defined"
+
+        # plane parametrized as (p - p0).n = 0
+        # line parametrized as p = l1 + l12*d with d Real
+        if l12 is None:
+            l12 = l2 - l1
+
+        if np.dot(l12,n) == 0:
+            return [0,0,0] # line is either in the plane or outside the plane
+        else:
+            d = np.dot((p0 - l1), n)/np.dot(l12,n)
+            return l1 + l12*d