ModuleAnalysisAndPlots

62def getETransmission(RayListIn, RayListOut) -> float:
63    """
64    Calculates the energy transmission from RayListIn to RayListOut in percent by summing up the
65    intensity-property of the individual rays.
66
67    Parameters
68    ----------
69        RayListIn, RayListOut : list(Ray)
70            Lists of objects of the ModuleOpticalRay.Ray-class.
71
72    Returns
73    -------
74        ETransmission : float
75    """
76    ETransmission = 100 * sum(Ray.intensity for Ray in RayListOut) / sum(Ray.intensity for Ray in RayListIn)
77    return ETransmission

 81def GetResultSummary(Detector, RayListAnalysed, verbose=False):
 82    """
 83    Calculate and return FocalSpotSize-standard-deviation and Duration-standard-deviation
 84    for the given Detector and RayList.
 85    If verbose, then also print a summary of the results for the given Detector.
 86
 87    Parameters
 88    ----------
 89        Detector : Detector
 90            An object of the ModuleDetector.Detector-class.
 91
 92        RayListAnalysed : list(Ray)
 93            List of objects of the ModuleOpticalRay.Ray-class.
 94
 95        verbose : bool
 96            Whether to print a result summary.
 97
 98    Returns
 99    -------
100        FocalSpotSizeSD : float
101
102        DurationSD : float
103    """
104    DetectorPointList2DCentre = Detector.get_PointList2DCentre(RayListAnalysed)
105    FocalSpotSizeSD = mp.StandardDeviation(DetectorPointList2DCentre)
106    DelayList = Detector.get_Delays(RayListAnalysed)
107    DurationSD = mp.StandardDeviation(DelayList)
108
109    if verbose:
110        FocalSpotSize = mgeo.DiameterPointList(DetectorPointList2DCentre)
111        summarystring = (
112            "At the detector distance of "
113            + "{:.3f}".format(Detector.get_distance())
114            + " mm we get:\n"
115            + "Spatial std : "
116            + "{:.3f}".format(FocalSpotSizeSD * 1e3)
117            + " \u03BCm and min-max: "
118            + "{:.3f}".format(FocalSpotSize * 1e3)
119            + " \u03BCm\n"
120            + "Temporal std : "
121            + "{:.3e}".format(DurationSD)
122            + " fs and min-max : "
123            + "{:.3e}".format(max(DelayList) - min(DelayList))
124            + " fs"
125        )
126
127        print(summarystring)
128
129    return FocalSpotSizeSD, DurationSD

133def SpotDiagram(RayListAnalysed, Detector, DrawAiryAndFourier=False, ColorCoded=None) -> plt.Figure:
134    """
135    Produce a an interactive figure with the spot diagram resulting from the RayListAnalysed
136    hitting the Detector.
137    The detector distance can be shifted with the left-right cursor keys.
138    If DrawAiryAndFourier is True, a circle with the Airy-spot-size will be shown.
139    The 'spots' can optionally be color-coded by specifying ColorCoded
140    as ["Intensity","Incidence","Delay"].
141
142    Parameters
143    ----------
144        RayListAnalysed : list(Ray)
145            List of objects of the ModuleOpticalRay.Ray-class.
146
147        Detector : Detector
148            An object of the ModuleDetector.Detector-class.
149
150        DrawAiryAndFourier : bool, optional
151            Whether to draw a circle with the Airy-spot-size. The default is False.
152
153        ColorCoded : str, optional
154            Color-code the spots according to one of ["Intensity","Incidence","Delay"]. The default is None.
155
156    Returns
157    -------
158        fig : matlplotlib-figure-handle.
159            Shows the interactive figure.
160    """
161    NumericalAperture = mp.ReturnNumericalAperture(RayListAnalysed, 1)  # NA determined from final ray bundle
162    Wavelength = RayListAnalysed[0].wavelength
163    if DrawAiryAndFourier:
164        AiryRadius = mp.ReturnAiryRadius(Wavelength, NumericalAperture) * 1e3  # in µm
165    else:
166        AiryRadius = 0
167
168    DectectorPoint2D_Xcoord, DectectorPoint2D_Ycoord, FocalSpotSize, SpotSizeSD = _getDetectorPoints(
169        RayListAnalysed, Detector
170    )
171
172    if ColorCoded == "Intensity":
173        IntensityList = [k.intensity for k in RayListAnalysed]
174        z = np.asarray(IntensityList)
175        zlabel = "Intensity (arb.u.)"
176        title = "Intensity + Spot Diagram\n press left/right to move detector position"
177        addLine = ""
178    elif ColorCoded == "Incidence":
179        IncidenceList = [np.rad2deg(k.incidence) for k in RayListAnalysed]  # degree
180        z = np.asarray(IncidenceList)
181        zlabel = "Incidence angle (deg)"
182        title = "Ray Incidence + Spot Diagram\n press left/right to move detector position"
183        addLine = ""
184    elif ColorCoded == "Delay":
185        DelayList = Detector.get_Delays(RayListAnalysed)
186        DurationSD = mp.StandardDeviation(DelayList)
187        z = np.asarray(DelayList)
188        zlabel = "Delay (fs)"
189        title = "Delay + Spot Diagram\n press left/right to move detector position"
190        addLine = "\n" + "{:.2f}".format(DurationSD) + " fs SD"
191    else:
192        z = "red"
193        title = "Spot Diagram\n press left/right to move detector position"
194        addLine = ""
195
196    Dist = Detector.get_distance()
197    distStep = min(50, max(0.0005, round(FocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000))  # in mm
198
199    plt.ion()
200    fig, ax = plt.subplots()
201    if DrawAiryAndFourier:
202        theta = np.linspace(0, 2 * np.pi, 100)
203        x = AiryRadius * np.cos(theta)
204        y = AiryRadius * np.sin(theta)  #
205        ax.plot(x, y, c="black")
206
207    foo = ax.scatter(
208        DectectorPoint2D_Xcoord,
209        DectectorPoint2D_Ycoord,
210        c=z,
211        s=15,
212        label="{:.3f}".format(Dist) + " mm\n" + "{:.1f}".format(SpotSizeSD * 1e3) + " \u03BCm SD" + addLine,
213    )
214
215    axisLim = 1.1 * max(AiryRadius, 0.5 * FocalSpotSize * 1000)
216    ax.set_xlim(-axisLim, axisLim)
217    ax.set_ylim(-axisLim, axisLim)
218
219    if ColorCoded == "Intensity" or ColorCoded == "Incidence" or ColorCoded == "Delay":
220        cbar = fig.colorbar(foo)
221        cbar.set_label(zlabel)
222
223    ax.legend(loc="upper right")
224    ax.set_xlabel("X (µm)")
225    ax.set_ylabel("Y (µm)")
226    ax.set_title(title)
227    # ax.margins(x=0)
228
229    movingDetector = Detector.copy_detector()
230
231    def press(event):
232        nonlocal Dist, distStep, movingDetector, ColorCoded, zlabel, cbar
233        if event.key == "right":
234            movingDetector.shiftByDistance(distStep)
235            Dist += distStep
236        elif event.key == "left":
237            if Dist > 1.5 * distStep:
238                movingDetector.shiftByDistance(-distStep)
239                Dist -= distStep
240            else:
241                movingDetector.shiftToDistance(0.5 * distStep)
242                Dist = 0.5 * distStep
243        else:
244            return None
245        newDectectorPoint2D_Xcoord, newDectectorPoint2D_Ycoord, newFocalSpotSize, newSpotSizeSD = _getDetectorPoints(
246            RayListAnalysed, movingDetector
247        )
248        xy = foo.get_offsets()
249        xy[:, 0] = newDectectorPoint2D_Xcoord
250        xy[:, 1] = newDectectorPoint2D_Ycoord
251        foo.set_offsets(xy)
252
253        if ColorCoded == "Delay":
254            newDelayList = np.asarray(movingDetector.get_Delays(RayListAnalysed))
255            newDurationSD = mp.StandardDeviation(newDelayList)
256            newaddLine = "\n" + "{:.2f}".format(newDurationSD) + " fs SD"
257            foo.set_array(newDelayList)
258            foo.set_clim(min(newDelayList), max(newDelayList))
259            cbar.update_normal(foo)
260
261        foo.set_label(
262            "{:.3f}".format(Dist) + " mm\n" + "{:.1f}".format(newSpotSizeSD * 1e3) + " \u03BCm SD" + newaddLine
263        )
264        ax.legend(loc="upper right")
265
266        axisLim = 1.1 * max(AiryRadius, 0.5 * newFocalSpotSize * 1000)
267        ax.set_xlim(-axisLim, axisLim)
268        ax.set_ylim(-axisLim, axisLim)
269
270        distStep = min(
271            50, max(0.0005, round(newFocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000)
272        )  # in mm
273
274        fig.canvas.draw_idle()
275
276    fig.canvas.mpl_connect("key_press_event", press)
277
278    plt.show()
279
280    return fig

359def DelayGraph(
360    RayListAnalysed, Detector, DeltaFT: (int, float), DrawAiryAndFourier=False, ColorCoded=None
361) -> plt.Figure:
362    """
363    Produce a an interactive figure with a spot diagram resulting from the RayListAnalysed
364    hitting the Detector, with the ray-delays shown in the 3rd dimension.
365    The detector distance can be shifted with the left-right cursor keys.
366    If DrawAiryAndFourier is True, a cylinder is shown whose diameter is the Airy-spot-size and
367    whose height is the Fourier-limited pulse duration 'given by 'DeltaFT'.
368    The 'spots' can optionally be color-coded by specifying ColorCoded as ["Intensity","Incidence"].
369
370    Parameters
371    ----------
372        RayListAnalysed : list(Ray)
373            List of objects of the ModuleOpticalRay.Ray-class.
374
375        Detector : Detector
376            An object of the ModuleDetector.Detector-class.
377
378        DeltaFT : (int, float)
379            The Fourier-limited pulse duration. Just used as a reference to compare the temporal spread
380            induced by the ray-delays.
381
382        DrawAiryAndFourier : bool, optional
383            Whether to draw a cylinder showing the Airy-spot-size and Fourier-limited-duration.
384            The default is False.
385
386        ColorCoded : str, optional
387            Color-code the spots according to one of ["Intensity","Incidence"].
388            The default is None.
389
390    Returns
391    -------
392        fig : matlplotlib-figure-handle.
393            Shows the interactive figure.
394    """
395    Dist = Detector.get_distance()
396    fig, NumericalAperture, AiryRadius, FocalSpotSize = _drawDelayGraph(
397        RayListAnalysed, Detector, Dist, DeltaFT, DrawAiryAndFourier, ColorCoded
398    )
399
400    distStep = min(50, max(0.0005, round(FocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000))  # in mm
401
402    movingDetector = Detector.copy_detector()
403
404    def press(event):
405        nonlocal Dist, distStep, movingDetector, fig
406        if event.key == "right":
407            movingDetector.shiftByDistance(distStep)
408            Dist += distStep
409            ax = fig.axes[0]
410            cam = [ax.azim, ax.elev, ax.dist]
411            fig, sameNumericalAperture, sameAiryRadius, newFocalSpotSize = _drawDelayGraph(
412                RayListAnalysed, movingDetector, Dist, DeltaFT, DrawAiryAndFourier, ColorCoded, fig
413            )
414            ax = fig.axes[0]
415            ax.azim, ax.elev, ax.dist = cam
416        elif event.key == "left":
417            if Dist > 1.5 * distStep:
418                movingDetector.shiftByDistance(-distStep)
419                Dist -= distStep
420            else:
421                movingDetector.shiftToDistance(0.5 * distStep)
422                Dist = 0.5 * distStep
423            ax = fig.axes[0]
424            cam = [ax.azim, ax.elev, ax.dist]
425
426            fig, sameNumericalAperture, sameAiryRadius, newFocalSpotSize = _drawDelayGraph(
427                RayListAnalysed, movingDetector, Dist, DeltaFT, DrawAiryAndFourier, ColorCoded, fig
428            )
429            ax = fig.axes[0]
430            ax.azim, ax.elev, ax.dist = cam
431        else:
432            return fig
433        distStep = min(
434            50, max(0.0005, round(newFocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000)
435        )  # in mm
436
437    fig.canvas.mpl_connect("key_press_event", press)
438
439    return fig

443def MirrorProjection(OpticalChain, ReflectionNumber: int, Detector=None, ColorCoded=None) -> plt.Figure:
444    """
445    Produce a plot of the ray impact points on the optical element with index 'ReflectionNumber'.
446    The points can be color-coded according ["Incidence","Intensity","Delay"], where the ray delay is
447    measured at the Detector.
448
449    Parameters
450    ----------
451        OpticalChain : OpticalChain
452           List of objects of the ModuleOpticalOpticalChain.OpticalChain-class.
453
454        ReflectionNumber : int
455            Index specifying the optical element on which you want to see the impact points.
456
457        Detector : Detector, optional
458            Object of the ModuleDetector.Detector-class. Only necessary to project delays. The default is None.
459
460        ColorCoded : str, optional
461            Specifies which ray property to color-code: ["Incidence","Intensity","Delay"]. The default is None.
462
463    Returns
464    -------
465        fig : matlplotlib-figure-handle.
466            Shows the figure.
467    """
468    from mpl_toolkits.axes_grid1 import make_axes_locatable
469
470    Position = OpticalChain.optical_elements[ReflectionNumber].position
471    n = OpticalChain.optical_elements[ReflectionNumber].normal
472    m = OpticalChain.optical_elements[ReflectionNumber].majoraxis
473
474    RayListAnalysed = OpticalChain.get_output_rays()[ReflectionNumber]
475    # transform rays into the mirror-support reference frame
476    # (same as mirror frame but without the shift by mirror-centre)
477    RayList = mgeo.TranslationRayList(RayListAnalysed, -Position)
478    RayList = mgeo.RotationRayList(RayList, n, np.array([0, 0, 1]))
479    mPrime = mgeo.RotationPoint(m, n, np.array([0, 0, 1]))
480    RayList = mgeo.RotationRayList(RayList, mPrime, np.array([1, 0, 0]))
481
482    x = np.asarray([k.point[0] for k in RayList])
483    y = np.asarray([k.point[1] for k in RayList])
484    if ColorCoded == "Intensity":
485        IntensityList = [k.intensity for k in RayListAnalysed]
486        z = np.asarray(IntensityList)
487        zlabel = "Intensity (arb.u.)"
488        title = "Ray intensity projected on mirror              "
489    elif ColorCoded == "Incidence":
490        IncidenceList = [np.rad2deg(k.incidence) for k in RayListAnalysed]  # in degree
491        z = np.asarray(IncidenceList)
492        zlabel = "Incidence angle (deg)"
493        title = "Ray incidence projected on mirror              "
494    elif ColorCoded == "Delay":
495        if Detector is not None:
496            z = np.asarray(Detector.get_Delays(RayListAnalysed))
497            zlabel = "Delay (fs)"
498            title = "Ray delay at detector projected on mirror              "
499        else:
500            raise ValueError("If you want to project ray delays, you must specify a detector.")
501    else:
502        z = "red"
503        title = "Ray impact points projected on mirror"
504
505    plt.ion()
506    fig = plt.figure()
507    ax = OpticalChain.optical_elements[ReflectionNumber].type.support._ContourSupport(fig)
508    p = plt.scatter(x, y, c=z, s=15)
509    if ColorCoded == "Delay" or ColorCoded == "Incidence" or ColorCoded == "Intensity":
510        divider = make_axes_locatable(ax)
511        cax = divider.append_axes("right", size="5%", pad=0.05)
512        cbar = fig.colorbar(p, cax=cax)
513        cbar.set_label(zlabel)
514    ax.set_xlabel("x (mm)")
515    ax.set_ylabel("y (mm)")
516    plt.title(title, loc="right")
517    plt.tight_layout()
518
519    bbox = ax.get_position()
520    bbox.set_points(bbox.get_points() - np.array([[0.01, 0], [0.01, 0]]))
521    ax.set_position(bbox)
522    plt.show()
523
524    return fig

528def RenderOpticalElement(OE, OEpoints=2000):
529    OpticPointList, edge_faces = OE.type.get_grid3D(OEpoints, edges=True)  # in the optic's coordinate system
530    # transform OpticPointList into "lab-frame"
531    OpticPointList = mgeo.TranslationPointList(OpticPointList, -OE.type.get_centre())
532    MirrorMajorAxisPrime = mgeo.RotationPoint(OE.majoraxis, OE.normal, np.array([0, 0, 1]))
533    OpticPointList = mgeo.RotationPointList(OpticPointList, np.array([1, 0, 0]), MirrorMajorAxisPrime)
534    OpticPointList = mgeo.RotationPointList(OpticPointList, np.array([0, 0, 1]), OE.normal)
535    OpticPointList = mgeo.TranslationPointList(OpticPointList, OE.position)
536
537    # plot 3D-drig of OE as little spheres
538    x = np.asarray([i[0] - OE.normal[0] * 0.5 for i in OpticPointList])
539    y = np.asarray([i[1] - OE.normal[1] * 0.5 for i in OpticPointList])
540    z = np.asarray([i[2] - OE.normal[2] * 0.5 for i in OpticPointList])
541
542    optic_pts = pv.PolyData(OpticPointList)
543
544    e = list(itertools.chain.from_iterable(edge_faces))
545    pts_coord = pv.PolyData(OpticPointList)
546    lines = list(
547        itertools.chain.from_iterable([[[2, e[i], e[i + 1]] for i in range(len(e) - 1)] for e in edge_faces])
548    )
549    faces = list(itertools.chain.from_iterable([[len(i) - 1] + i[:-1] for i in edge_faces]))
550    if lines == []:
551        lines = [0]
552    if faces == []:
553        faces = [0]
554    edges = pv.PolyData(
555        OpticPointList,
556        lines=lines,
557        faces=faces,
558    )
559    # Can't figure out some edge cases such as when part of the support is outside of the mirror
560    tess = pts_coord.delaunay_2d(edge_source=edges)
561    triangles = tess.faces.reshape(-1, 4)[:, 1:]
562    return optic_pts, tess

564def RenderRays(RayListHistory, EndDistance=None, maxRays=150, color_by_number = True):
565    meshes = []
566    # Ray display
567    for k in range(len(RayListHistory)):
568        x = []
569        y = []
570        z = []
571        connections = []
572        if k != len(RayListHistory) - 1:
573            knums = list(
574                map(lambda x: x.number, RayListHistory[k])
575            )  # make a list of all ray numbers that are still in the game
576            if len(RayListHistory[k + 1]) > maxRays:
577                rays_to_render = np.random.choice(RayListHistory[k + 1], maxRays, replace=False)
578            else:
579                rays_to_render = RayListHistory[k + 1]
580
581            for j in rays_to_render:
582                indx = knums.index(j.number)
583                i = RayListHistory[k][indx]
584                Point1 = i.point
585                Point2 = j.point
586                x += [Point1[0], Point2[0]]
587                y += [Point1[1], Point2[1]]
588                z += [Point1[2], Point2[2]]
589
590        else:
591            if len(RayListHistory[k]) > maxRays:
592                rays_to_render = np.random.choice(RayListHistory[k], maxRays, replace=False)
593            else:
594                rays_to_render = RayListHistory[k]
595
596            for j in rays_to_render:
597                Point = j.point
598                Vector = j.vector
599                x += [Point[0], Point[0] + Vector[0] * EndDistance]
600                y += [Point[1], Point[1] + Vector[1] * EndDistance]
601                z += [Point[2], Point[2] + Vector[2] * EndDistance]
602        points = np.column_stack((x, y, z))
603        meshes += [pv.line_segments_from_points(points)]
604    return meshes

606def generate_distinct_colors(num_colors):
607    # Get a color palette from colorcet
608    palette = cc.glasbey
609
610    # Make sure the number of colors does not exceed the palette length
611    num_colors = min(num_colors, len(palette))
612
613    # Slice the palette to get the desired number of colors
614    distinct_colors = palette[:num_colors]
615
616    return distinct_colors

618def RayRenderGraph(OpticalChain, EndDistance=None, maxRays=150, OEpoints=2000, scale_spheres=0.0):
619    """
620    Renders an image of the Optical setup and the traced rays.
621
622    Parameters
623    ----------
624        OpticalChain : OpticalChain
625            List of objects of the ModuleOpticalOpticalChain.OpticalChain-class.
626
627        EndDistance : float, optional
628            The rays of the last ray bundle are drawn with a length given by EndDistance (in mm). If not specified,
629            this distance is set to that between the source point and the 1st optical element.
630
631        maxRays: int
632            The maximum number of rays to render. Rendering all the traced rays is a insufferable resource hog
633            and not required for a nice image. Default is 150.
634
635        OEpoints : int
636            How many little spheres to draw to represent the optical elements.  Default is 2000.
637
638    Returns
639    -------
640        fig : Pyvista-figure-handle.
641            Shows the figure.
642    """
643
644    RayListHistory = [OpticalChain.source_rays] + OpticalChain.get_output_rays()
645
646    if EndDistance is None:
647        EndDistance = np.linalg.norm(OpticalChain.source_rays[0].point - OpticalChain.optical_elements[0].position)
648
649    print("...rendering image of optical chain...", end="", flush=True)
650    fig = pv.Plotter(window_size=(1500, 500), notebook=False)
651
652    ray_meshes = RenderRays(RayListHistory, EndDistance, maxRays)
653    colors = generate_distinct_colors(len(ray_meshes))
654    for i,ray in enumerate(ray_meshes):
655        color = pv.Color(colors[i])
656        fig.add_mesh(ray, color=color)
657
658    # Optics display
659    for i,OE in enumerate(OpticalChain.optical_elements):
660        pointcloud, mesh = RenderOpticalElement(OE, OEpoints)
661        color = pv.Color(colors[i+1])
662        color = colorsys.hsv_to_rgb(*(colorsys.rgb_to_hsv(*color[:-1]))*np.array([1,0.2,1]))
663        fig.add_mesh(mesh, color = color)
664        fig.add_mesh(pointcloud, point_size=scale_spheres, color = color, render_points_as_spheres = True)
665    fig.show()
666    print(
667        "\r\033[K", end="", flush=True
668    )  # move to beginning of the line with \r and then delete the whole line with \033[K
669    return fig

674def show():
675    plt.show(block=False)