Skip to content
Snippets Groups Projects
GeoBeams.py 12 KiB
Newer Older
# -*- coding: utf-8 -*-
""" Geometric beam calculator for SCS.

    Copyright (2021) SCS Team.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
from sympy.physics.optics import GeometricRay, FreeSpace, ThinLens
from sympy.solvers import solve
from sympy import symbols, 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
Loïc Le Guyader's avatar
Loïc Le Guyader committed
            'fBOZ_x': 0.25,  # BOZ horizontal focal length
            'fBOZ_y': 0.25,  # BOZ vertical focal distance
            'theta_grating': 0,  # grating deflection angle in rad
            'dSAMZ': 0.0,  # Sample position
            'ddetZ': 2.0,  # Detector 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')
Loïc Le Guyader's avatar
Loïc Le Guyader committed
        fVFM, fHFM, fBOZ_x, fBOZ_y = symbols('fVFM, fHFM, fBOZ_x, fBOZ_y')
        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)
Loïc Le Guyader's avatar
Loïc Le Guyader committed
        self.HBOZ = simplify(ThinLens(fBOZ_x)*self.HFM)
        self.VBOZ = simplify(ThinLens(fBOZ_y)*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
Loïc Le Guyader's avatar
Loïc Le Guyader committed
        ray = (ThinLens(self.elems['fBOZ_x']) *
               FreeSpace(-self.elems['dHFM'] + self.elems['dBOZ'])*ray)
        z_x += [self.elems['dBOZ']]
        x += [ray[0].evalf()]

        # detector
        ray = FreeSpace(-self.elems['dBOZ']+self.elems['ddetZ'])*ray
        z_x += [self.elems['ddetZ']]
        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
Loïc Le Guyader's avatar
Loïc Le Guyader committed
        ray = (ThinLens(self.elems['fBOZ_y']) *
               FreeSpace(-self.elems['dVFM']+self.elems['dBOZ'])*ray)
        z_y += [self.elems['dBOZ']]
        y += [ray[0].evalf()]

        # detector
        ray = FreeSpace(-self.elems['dBOZ']+self.elems['ddetZ'])*ray
        z_y += [self.elems['ddetZ']]
    def plot(self, conf1=None, conf2=None):
        fig, ax = plt.subplots(2, 1, figsize=(6, 4), sharex=True)
        ax0in = ax[0].inset_axes([0.1, 0.65, 0.2, 0.25])
        ax1in = ax[1].inset_axes([0.1, 0.65, 0.2, 0.25])
        axs = [ax, ax0in, ax1in]

        if conf1 is not None:
            self.elems['fBOZ_x'] = conf1['F_x']
            self.elems['fBOZ_y'] = conf1['F_y']
            self.elems['theta_grating'] = conf1['theta_grating']
        # plot either a given conf or whatever the last calculation was
        c_rr = hsv_to_rgb([10/360, 100/100, 100/100])
        self._plot(axs, c_rr)


        if conf2 is not None:
            self.elems['fBOZ_x'] = conf2['F_x']
            self.elems['fBOZ_y'] = conf2['F_y']
            self.elems['theta_grating'] = conf2['theta_grating']
            c_rb = hsv_to_rgb([220/360, 100/100, 100/100])
            self._plot(axs, c_rb)

        ax[0].axvline(self.elems['dSAMZ'], ls='--', c='k', alpha=0.5)
        ax0in.axvline(self.elems['dSAMZ'], ls='--', c='k', alpha=0.5)
        ax[1].axvline(self.elems['dSAMZ'], ls='--', c='k', alpha=0.5)
        ax1in.axvline(self.elems['dSAMZ'], ls='--', c='k', alpha=0.5)

        ax0in.set_xlim([-0.1, 0.3])
        ax0in.set_ylim([-0.3, 0.3])
        ax1in.set_xlim([-0.1, 0.3])
        ax1in.set_ylim([-0.3, 0.3])

        ax[0].set_ylabel('x (mm)')
        ax[1].set_ylabel('y (mm)')
        ax[1].set_xlabel('z (m)')

        return fig

    def _plot(self, axs, color):
        ax, ax0in, ax1in = axs

        xmin, xmax = [np.Inf, -np.Inf]
        ymin, ymax = [np.Inf, -np.Inf]
        xs = np.empty((4, 4))
        ys = np.empty((4, 4))
        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 kk, (x1, y1) in enumerate(zip(
                    [self.elems['IHFw']/2, -self.elems['IHFw']/2],
                    [self.elems['EXw']/2, -self.elems['EXw']/2])):

                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()

                xs[2*k + kk, :] = 1e3*np.array(x)
                ys[2*k + kk, :] = 1e3*np.array(y)

                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]

        # fill_between to show beams
        for k in range(4):
            for kk in range(k+1, 4):
                ax[0].fill_between(z_x, xs[k, :], xs[kk, :],
                                   color=color, alpha=0.2, linewidth=0)
                ax0in.fill_between(z_x, xs[k, :], xs[kk, :],
                                   color=color, alpha=0.2, linewidth=0)
                ax[1].fill_between(z_y, ys[k, :], ys[kk, :],
                                   color=color, alpha=0.2, linewidth=0)
                ax1in.fill_between(z_y, ys[k, :], ys[kk, :],
                                   color=color, alpha=0.2, linewidth=0)

    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 kk, (x1, y1) in enumerate(zip(
                    [self.elems['IHFw']/2, -self.elems['IHFw']/2],
                    [self.elems['EXw']/2, -self.elems['EXw']/2])):

                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 of it
        else:
            d = np.dot((p0 - l1), n)/np.dot(l12, n)
            return l1 + l12*d