ModuleAnalysisAndPlots

  1"""
  2Provides functions for analysis of the output ray-bundles calculated through ray-tracing, and also for the ART's standard visualization options.
  3
  4Implements the following functions:
  5    - DrawSpotDiagram: Produces an interactive figure with the spot diagram on the selected Detector.
  6    - DrawDelaySpots: Produces a an interactive figure with a spot diagram resulting from the RayListAnalysed
  7    - _drawDelayGraph: Auxiliary function for DelayGraph.
  8    - DrawMirrorProjection: Produce a plot of the ray impact points on the optical element with index 'ReflectionNumber'.
  9    - DrawSetup: Renders an image of the Optical setup and the traced rays.
 10    - DrawAsphericity: Displays a map of the asphericity of the mirror.
 11    - DrawCaustics: Displays the caustics of the rays on the detector.
 12
 13It also implements the following methods:
 14    -OpticalChain:
 15        - render: Method to render the optical setup.
 16        - drawSpotDiagram: Method to draw the spot diagram on the selected detector.
 17        - drawDelaySpots: Method to draw the delay spots on the detector.
 18        - drawMirrorProjection: Method to draw the mirror projection.
 19        - drawCaustics: Method to draw the caustics on the detector.
 20    -Mirror:
 21        - drawAsphericity: Method to draw the asphericity of the mirror.
 22
 23
 24Created in Apr 2020
 25
 26@author: Anthony Guillaume + Stefan Haessler + Andre Kalouguine
 27"""
 28# %% Module imports
 29import numpy as np
 30import matplotlib.pyplot as plt
 31from mpl_toolkits.mplot3d import Axes3D
 32import pyvista as pv
 33import pyvistaqt as pvqt
 34import colorcet as cc
 35from colorsys import rgb_to_hls, hls_to_rgb
 36import logging
 37
 38logger = logging.getLogger(__name__)
 39
 40
 41import ARTcore.ModuleProcessing as mp
 42import ARTcore.ModuleGeometry as mgeo
 43import ARTcore.ModuleDetector as mdet
 44import ARTcore.ModuleMirror as mmirror
 45import ART.ModulePlottingMethods as mpm
 46import ART.ModulePlottingUtilities as mpu
 47from ART.ModulePlottingUtilities import Observable
 48
 49import ARTcore.ModuleOpticalChain as moc
 50import ART.ModuleAnalysis as man
 51import itertools
 52from copy import copy
 53
 54
 55# %% Spot diagram on detector
 56def DrawSpotDiagram(OpticalChain, 
 57                Detector = "Focus",
 58                DrawAiryAndFourier=False, 
 59                DrawFocalContour=False,
 60                DrawFocal=False,
 61                ColorCoded=None,
 62                Observer = None) -> plt.Figure:
 63    """
 64    Produce an interactive figure with the spot diagram on the selected Detector.
 65    The detector distance can be shifted with the left-right cursor keys. Doing so will actually move the detector.
 66    If DrawAiryAndFourier is True, a circle with the Airy-spot-size will be shown.
 67    If DrawFocalContour is True, the focal contour calculated from some of the rays will be shown.
 68    If DrawFocal is True, a heatmap calculated from some of the rays will be shown. 
 69    The 'spots' can optionally be color-coded by specifying ColorCoded, which can be one of ["Intensity","Incidence","Delay"].
 70
 71    Parameters
 72    ----------
 73        RayListAnalysed : list(Ray)
 74            List of objects of the ModuleOpticalRay.Ray-class.
 75
 76        Detector : Detector or str, optional
 77            An object of the ModuleDetector.Detector-class or the name of the detector. The default is "Focus".
 78
 79        DrawAiryAndFourier : bool, optional
 80            Whether to draw a circle with the Airy-spot-size. The default is False.
 81        
 82        DrawFocalContour : bool, optional
 83            Whether to draw the focal contour. The default is False.
 84
 85        DrawFocal : bool, optional
 86            Whether to draw the focal heatmap. The default is False.
 87
 88        ColorCoded : str, optional
 89            Color-code the spots according to one of ["Intensity","Incidence","Delay"]. The default is None.
 90
 91        Observer : Observer, optional
 92            An observer object. If none, then we just create a copy of the detector and move it when pressing left-right. 
 93            However, if an observer is specified, then we will change the value of the observer and it will issue 
 94            the required callbacks to update several plots at the same time.
 95
 96    Returns
 97    -------
 98        fig : matlplotlib-figure-handle.
 99            Shows the interactive figure.
100    """
101    if isinstance(Detector, str):
102        Detector = OpticalChain.detectors[Detector]
103    Index = Detector.index
104    movingDetector = copy(Detector) # We will move this detector when pressing left-right
105    if Observer is None:
106        detectorPosition = Observable(movingDetector.distance) # We will observe the distance of this detector
107    else:
108        detectorPosition = Observer
109        movingDetector.distance = detectorPosition.value
110
111    detectorPosition.register_calculation(lambda x: movingDetector.set_distance(x))
112
113    RayListAnalysed = OpticalChain.get_output_rays()[Index]
114
115    NumericalAperture = man.GetNumericalAperture(RayListAnalysed, 1)  # NA determined from final ray bundle
116    MaxWavelength = np.max([i.wavelength for i in RayListAnalysed])
117    if DrawAiryAndFourier:
118        AiryRadius = man.GetAiryRadius(MaxWavelength, NumericalAperture) * 1e3  # in µm
119    else:
120        AiryRadius = 0
121    
122    if DrawFocalContour or DrawFocal:
123        X,Y,Z = man.GetDiffractionFocus(OpticalChain, movingDetector, Index)
124        Z/=np.max(Z) 
125
126    DectectorPoint2D_Xcoord, DectectorPoint2D_Ycoord, FocalSpotSize, SpotSizeSD = mpu._getDetectorPoints(
127        RayListAnalysed, movingDetector
128    )
129
130    match ColorCoded:
131        case "Intensity":
132            IntensityList = [k.intensity for k in RayListAnalysed]
133            z = np.asarray(IntensityList)
134            zlabel = "Intensity (arb.u.)"
135            title = "Intensity + Spot Diagram\n press left/right to move detector position"
136            addLine = ""
137        case "Incidence":
138            IncidenceList = [np.rad2deg(k.incidence) for k in RayListAnalysed]  # degree
139            z = np.asarray(IncidenceList)
140            zlabel = "Incidence angle (deg)"
141            title = "Ray Incidence + Spot Diagram\n press left/right to move detector position"
142            addLine = ""
143        case "Delay":
144            DelayList = movingDetector.get_Delays(RayListAnalysed)
145            DurationSD = mp.StandardDeviation(DelayList)
146            z = np.asarray(DelayList)
147            zlabel = "Delay (fs)"
148            title = "Delay + Spot Diagram\n press left/right to move detector position"
149            addLine = "\n" + "{:.2f}".format(DurationSD) + " fs SD"
150        case _:
151            z = "red"
152            title = "Spot Diagram\n press left/right to move detector position"
153            addLine = ""
154
155    distStep = min(50, max(0.0005, round(FocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000))  # in mm
156
157    plt.ion()
158    fig, ax = plt.subplots()
159    if DrawFocal:
160        focal = ax.pcolormesh(X*1e3,Y*1e3,Z)
161    if DrawFocalContour:
162        levels = [1/np.e**2, 0.5]
163        contour = ax.contourf(X*1e3, Y*1e3, Z, levels=levels, cmap='gray')
164
165    if DrawAiryAndFourier:
166        theta = np.linspace(0, 2 * np.pi, 100)
167        x = AiryRadius * np.cos(theta)
168        y = AiryRadius * np.sin(theta)  #
169        ax.plot(x, y, c="black")
170        
171
172    foo = ax.scatter(
173        DectectorPoint2D_Xcoord,
174        DectectorPoint2D_Ycoord,
175        c=z,
176        s=15,
177        label="{:.3f}".format(detectorPosition.value) + " mm\n" + "{:.1f}".format(SpotSizeSD * 1e3) + " \u03BCm SD" + addLine,
178    )
179
180    axisLim = 1.1 * max(AiryRadius, 0.5 * FocalSpotSize * 1000)
181    ax.set_xlim(-axisLim, axisLim)
182    ax.set_ylim(-axisLim, axisLim)
183
184    if ColorCoded == "Intensity" or ColorCoded == "Incidence" or ColorCoded == "Delay":
185        cbar = fig.colorbar(foo)
186        cbar.set_label(zlabel)
187
188    ax.legend(loc="upper right")
189    ax.set_xlabel("X (µm)")
190    ax.set_ylabel("Y (µm)")
191    ax.set_title(title)
192    # ax.margins(x=0)
193
194
195    def update_plot(new_value):
196        nonlocal movingDetector, ColorCoded, zlabel, cbar, detectorPosition, foo, distStep, focal, contour, levels, Index, RayListAnalysed
197
198        newDectectorPoint2D_Xcoord, newDectectorPoint2D_Ycoord, newFocalSpotSize, newSpotSizeSD = mpu._getDetectorPoints(
199            RayListAnalysed, movingDetector
200        )
201
202        if DrawFocal:
203            focal.set_array(Z)
204        if DrawFocalContour:
205            levels = [1/np.e**2, 0.5]
206            for coll in contour.collections:
207                coll.remove()  # Remove old contour lines
208            contour = ax.contourf(X * 1e3, Y * 1e3, Z, levels=levels, cmap='gray')
209        
210        xy = foo.get_offsets()
211        xy[:, 0] = newDectectorPoint2D_Xcoord
212        xy[:, 1] = newDectectorPoint2D_Ycoord
213        foo.set_offsets(xy)
214
215
216        if ColorCoded == "Delay":
217            newDelayList = np.asarray(movingDetector.get_Delays(RayListAnalysed))
218            newDurationSD = mp.StandardDeviation(newDelayList)
219            newaddLine = "\n" + "{:.2f}".format(newDurationSD) + " fs SD"
220            foo.set_array(newDelayList)
221            foo.set_clim(min(newDelayList), max(newDelayList))
222            cbar.update_normal(foo)
223        else:
224            newaddLine = ""
225
226        foo.set_label(
227            "{:.3f}".format(detectorPosition.value) + " mm\n" + "{:.1f}".format(newSpotSizeSD * 1e3) + " \u03BCm SD" + newaddLine
228        )
229        ax.legend(loc="upper right")
230
231        axisLim = 1.1 * max(AiryRadius, 0.5 * newFocalSpotSize * 1000)
232        ax.set_xlim(-axisLim, axisLim)
233        ax.set_ylim(-axisLim, axisLim)
234
235        distStep = min(
236            50, max(0.0005, round(newFocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000)
237        )  # in mm
238
239        fig.canvas.draw_idle()
240
241
242    def press(event):
243        nonlocal detectorPosition, distStep
244        if event.key == "right":
245            detectorPosition.value += distStep
246        elif event.key == "left":
247            if detectorPosition.value > 1.5 * distStep:
248                detectorPosition.value -= distStep
249            else:
250                detectorPosition.value = 0.5 * distStep
251        else:
252            return None
253
254    fig.canvas.mpl_connect("key_press_event", press)
255
256    plt.show()
257
258    detectorPosition.register(update_plot)
259
260
261    return fig, detectorPosition
262
263moc.OpticalChain.drawSpotDiagram = DrawSpotDiagram
264
265# %% Delay graph on detector (3D spot diagram)
266def DrawDelaySpots(OpticalChain, 
267                DeltaFT: tuple[int, float],
268                Detector = "Focus",
269                DrawAiryAndFourier=False, 
270                ColorCoded=None,
271                Observer = None
272                ) -> plt.Figure:
273    """
274    Produce a an interactive figure with a spot diagram resulting from the RayListAnalysed
275    hitting the Detector, with the ray-delays shown in the 3rd dimension.
276    The detector distance can be shifted with the left-right cursor keys.
277    If DrawAiryAndFourier is True, a cylinder is shown whose diameter is the Airy-spot-size and
278    whose height is the Fourier-limited pulse duration 'given by 'DeltaFT'.
279    
280    The 'spots' can optionally be color-coded by specifying ColorCoded as ["Intensity","Incidence"].
281
282    Parameters
283    ----------
284        RayListAnalysed : list(Ray)
285            List of objects of the ModuleOpticalRay.Ray-class.
286
287        Detector : Detector
288            An object of the ModuleDetector.Detector-class.
289
290        DeltaFT : (int, float)
291            The Fourier-limited pulse duration. Just used as a reference to compare the temporal spread
292            induced by the ray-delays.
293
294        DrawAiryAndFourier : bool, optional
295            Whether to draw a cylinder showing the Airy-spot-size and Fourier-limited-duration.
296            The default is False.
297
298        ColorCoded : str, optional
299            Color-code the spots according to one of ["Intensity","Incidence"].
300            The default is None.
301
302    Returns
303    -------
304        fig : matlplotlib-figure-handle.
305            Shows the interactive figure.
306    """
307    if isinstance(Detector, str):
308        Det = OpticalChain.detectors[Detector]
309    else:
310        Det = Detector
311    Index = Det.index
312    Detector = copy(Det)
313    if Observer is None:
314        detectorPosition = Observable(Detector.distance)
315    else:
316        detectorPosition = Observer
317        Detector.distance = detectorPosition.value
318    
319    RayListAnalysed = OpticalChain.get_output_rays()[Index]
320    fig, NumericalAperture, AiryRadius, FocalSpotSize = _drawDelayGraph(
321        RayListAnalysed, Detector, detectorPosition.value, DeltaFT, DrawAiryAndFourier, ColorCoded
322    )
323
324    distStep = min(50, max(0.0005, round(FocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000))  # in mm
325
326    movingDetector = copy(Detector)
327
328    def update_plot(new_value):
329        nonlocal movingDetector, ColorCoded, detectorPosition, distStep, fig
330        ax = fig.axes[0]
331        cam = [ax.azim, ax.elev, ax._dist]
332        fig, sameNumericalAperture, sameAiryRadius, newFocalSpotSize = _drawDelayGraph(
333            RayListAnalysed, movingDetector, detectorPosition.value, DeltaFT, DrawAiryAndFourier, ColorCoded, fig
334        )
335        ax = fig.axes[0]
336        ax.azim, ax.elev, ax._dist = cam
337        distStep = min(
338            50, max(0.0005, round(newFocalSpotSize / 8 / np.arcsin(NumericalAperture) * 10000) / 10000)
339        )
340        return fig
341
342    def press(event):
343        nonlocal detectorPosition, distStep, movingDetector, fig
344        if event.key == "right":
345            detectorPosition.value += distStep
346        elif event.key == "left":
347            if detectorPosition.value > 1.5 * distStep:
348                detectorPosition.value -= distStep
349            else:
350                detectorPosition.value = 0.5 * distStep
351
352    fig.canvas.mpl_connect("key_press_event", press)
353    detectorPosition.register(update_plot)
354    detectorPosition.register_calculation(lambda x: movingDetector.set_distance(x))
355
356    return fig, Observable
357
358
359def _drawDelayGraph(RayListAnalysed, Detector, Distance, DeltaFT, DrawAiryAndFourier=False, ColorCoded=None, fig=None):
360    """
361    Draws the 3D-delay-spot-diagram for a fixed detector. See more doc in the function below.
362    """
363    NumericalAperture = man.GetNumericalAperture(RayListAnalysed, 1)  # NA determined from final ray bundle
364    Wavelength = RayListAnalysed[0].wavelength
365    AiryRadius = man.GetAiryRadius(Wavelength, NumericalAperture) * 1e3  # in µm
366
367    DectectorPoint2D_Xcoord, DectectorPoint2D_Ycoord, FocalSpotSize, SpotSizeSD = mpu._getDetectorPoints(
368        RayListAnalysed, Detector
369    )
370
371    DelayList = Detector.get_Delays(RayListAnalysed)
372    DurationSD = mp.StandardDeviation(DelayList)
373
374    if ColorCoded == "Intensity":
375        IntensityList = [k.intensity for k in RayListAnalysed]
376    elif ColorCoded == "Incidence":
377        IncidenceList = [np.rad2deg(k.incidence) for k in RayListAnalysed]  # in degree
378
379    plt.ion()
380    if fig is None:
381        fig = plt.figure()
382    else:
383        fig.clear(keep_observers=True)
384
385    ax = Axes3D(fig)
386    fig.add_axes(ax)
387    ax.set_xlabel("X (µm)")
388    ax.set_ylabel("Y (µm)")
389    ax.set_zlabel("Delay (fs)")
390
391    labelLine = (
392        "{:.3f}".format(Distance)
393        + " mm\n"
394        + "{:.1f}".format(SpotSizeSD * 1e3)
395        + " \u03BCm SD\n"
396        + "{:.2f}".format(DurationSD)
397        + " fs SD"
398    )
399    if ColorCoded == "Intensity":
400        ax.scatter(DectectorPoint2D_Xcoord, DectectorPoint2D_Ycoord, DelayList, s=4, c=IntensityList, label=labelLine)
401        # plt.title('Delay + Intensity graph\n press left/right to move detector position')
402        fig.suptitle("Delay + Intensity graph\n press left/right to move detector position")
403
404    elif ColorCoded == "Incidence":
405        ax.scatter(DectectorPoint2D_Xcoord, DectectorPoint2D_Ycoord, DelayList, s=4, c=IncidenceList, label=labelLine)
406        # plt.title('Delay + Incidence graph\n press left/right to move detector position')
407        fig.suptitle("Delay + Incidence graph\n press left/right to move detector position")
408    else:
409        ax.scatter(DectectorPoint2D_Xcoord, DectectorPoint2D_Ycoord, DelayList, s=4, c=DelayList, label=labelLine)
410        # plt.title('Delay graph\n press left/right to move detector position')
411        fig.suptitle("Delay graph\n press left/right to move detector position")
412
413    ax.legend(loc="upper right")
414
415    if DrawAiryAndFourier:
416        x = np.linspace(-AiryRadius, AiryRadius, 40)
417        z = np.linspace(np.mean(DelayList) - DeltaFT * 0.5, np.mean(DelayList) + DeltaFT * 0.5, 40)
418        x, z = np.meshgrid(x, z)
419        y = np.sqrt(AiryRadius**2 - x**2)
420        ax.plot_wireframe(x, y, z, color="grey", alpha=0.1)
421        ax.plot_wireframe(x, -y, z, color="grey", alpha=0.1)
422
423    axisLim = 1.1 * max(AiryRadius, 0.5 * FocalSpotSize * 1000)
424    ax.set_xlim(-axisLim, axisLim)
425    ax.set_ylim(-axisLim, axisLim)
426    #ax.set_zlim(-axisLim/3*10, axisLim/3*10) #same scaling as spatial axes 
427    ax.set_zlim(-DeltaFT, DeltaFT)
428    
429    plt.show()
430    fig.canvas.draw()
431
432    return fig, NumericalAperture, AiryRadius, FocalSpotSize
433
434moc.OpticalChain.drawDelaySpots = DrawDelaySpots
435
436# %% Mirror projection
437def DrawMirrorProjection(OpticalChain, ReflectionNumber: int, ColorCoded=None, Detector="") -> plt.Figure:
438    """
439    Produce a plot of the ray impact points on the optical element with index 'ReflectionNumber'.
440    The points can be color-coded according ["Incidence","Intensity","Delay"], where the ray delay is
441    measured at the Detector.
442
443    Parameters
444    ----------
445        OpticalChain : OpticalChain
446           List of objects of the ModuleOpticalOpticalChain.OpticalChain-class.
447
448        ReflectionNumber : int
449            Index specifying the optical element on which you want to see the impact points.
450
451        Detector : Detector, optional
452            Object of the ModuleDetector.Detector-class. Only necessary to project delays. The default is None.
453
454        ColorCoded : str, optional
455            Specifies which ray property to color-code: ["Incidence","Intensity","Delay"]. The default is None.
456
457    Returns
458    -------
459        fig : matlplotlib-figure-handle.
460            Shows the figure.
461    """
462    from mpl_toolkits.axes_grid1 import make_axes_locatable
463    if isinstance(Detector, str):
464        if Detector == "":
465            Detector = None
466        else:
467            Detector = OpticalChain.detectors[Detector]
468
469    Position = OpticalChain[ReflectionNumber].position
470    q = OpticalChain[ReflectionNumber].orientation
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    r0 = OpticalChain[ReflectionNumber].r0
478    RayList = [r.to_basis(*OpticalChain[ReflectionNumber].basis) for r in RayListAnalysed]
479
480    x = np.asarray([k.point[0] for k in RayList]) - r0[0]
481    y = np.asarray([k.point[1] for k in RayList]) - r0[1]
482    if ColorCoded == "Intensity":
483        IntensityList = [k.intensity for k in RayListAnalysed]
484        z = np.asarray(IntensityList)
485        zlabel = "Intensity (arb.u.)"
486        title = "Ray intensity projected on mirror              "
487    elif ColorCoded == "Incidence":
488        IncidenceList = [np.rad2deg(k.incidence) for k in RayListAnalysed]  # in degree
489        z = np.asarray(IncidenceList)
490        zlabel = "Incidence angle (deg)"
491        title = "Ray incidence projected on mirror              "
492    elif ColorCoded == "Delay":
493        if Detector is not None:
494            z = np.asarray(Detector.get_Delays(RayListAnalysed))
495            zlabel = "Delay (fs)"
496            title = "Ray delay at detector projected on mirror              "
497        else:
498            raise ValueError("If you want to project ray delays, you must specify a detector.")
499    else:
500        z = "red"
501        title = "Ray impact points projected on mirror"
502
503    plt.ion()
504    fig = plt.figure()
505    ax = OpticalChain.optical_elements[ReflectionNumber].support._ContourSupport(fig)
506    p = plt.scatter(x, y, c=z, s=15)
507    if ColorCoded == "Delay" or ColorCoded == "Incidence" or ColorCoded == "Intensity":
508        divider = make_axes_locatable(ax)
509        cax = divider.append_axes("right", size="5%", pad=0.05)
510        cbar = fig.colorbar(p, cax=cax)
511        cbar.set_label(zlabel)
512    ax.set_xlabel("x (mm)")
513    ax.set_ylabel("y (mm)")
514    plt.title(title, loc="right")
515    plt.tight_layout()
516
517    bbox = ax.get_position()
518    bbox.set_points(bbox.get_points() - np.array([[0.01, 0], [0.01, 0]]))
519    ax.set_position(bbox)
520    plt.show()
521
522    return fig
523
524
525moc.OpticalChain.drawMirrorProjection = DrawMirrorProjection
526
527# %% Setup rendering
528def DrawSetup(OpticalChain, 
529                   EndDistance=None, 
530                   maxRays=300, 
531                   OEpoints=2000, 
532                   draw_mesh=False, 
533                   cycle_ray_colors = False,
534                   impact_points = False,
535                   DrawDetectors=True,
536                   DetectedRays = False,
537                   Observers = dict()):
538    """
539    Renders an image of the Optical setup and the traced rays.
540
541    Parameters
542    ----------
543        OpticalChain : OpticalChain
544            List of objects of the ModuleOpticalOpticalChain.OpticalChain-class.
545
546        EndDistance : float, optional
547            The rays of the last ray bundle are drawn with a length given by EndDistance (in mm). If not specified,
548            this distance is set to that between the source point and the 1st optical element.
549
550        maxRays: int
551            The maximum number of rays to render. Rendering all the traced rays is a insufferable resource hog
552            and not required for a nice image. Default is 150.
553
554        OEpoints : int
555            How many little spheres to draw to represent the optical elements.  Default is 2000.
556
557    Returns
558    -------
559        fig : Pyvista-figure-handle.
560            Shows the figure.
561    """
562
563    RayListHistory = [OpticalChain.source_rays] + OpticalChain.get_output_rays()
564
565    if EndDistance is None:
566        EndDistance = np.linalg.norm(OpticalChain.source_rays[0].point - OpticalChain.optical_elements[0].position)
567
568    print("...rendering image of optical chain...", end="", flush=True)
569    fig = pvqt.BackgroundPlotter(window_size=(1500, 500), notebook=False) # Opening a window
570    fig.set_background('white')
571    
572    if cycle_ray_colors:
573        colors = mpu.generate_distinct_colors(len(OpticalChain)+1)
574    else:
575        colors = [[0.7, 0, 0]]*(len(OpticalChain)+1) # Default color: dark red
576
577    # Optics display
578    # For each optic we will send the figure to the function _RenderOpticalElement and it will add the optic to the figure
579    for i,OE in enumerate(OpticalChain.optical_elements):
580        color = pv.Color(colors[i+1])
581        rgb = color.float_rgb
582        h, l, s = rgb_to_hls(*rgb)
583        s = max(0, min(1, s * 0.3))  # Decrease saturation
584        l = max(0, min(1, l + 0.1))  # Increase lightness
585        new_rgb = hls_to_rgb(h, l, s)
586        darkened_color = pv.Color(new_rgb)
587        mpm._RenderOpticalElement(fig, OE, OEpoints, draw_mesh, darkened_color, index=i)
588    ray_meshes = mpm._RenderRays(RayListHistory, EndDistance, maxRays)
589    for i,ray in enumerate(ray_meshes):
590        color = pv.Color(colors[i])
591        fig.add_mesh(ray, color=color, name=f"RayBundle_{i}")
592    if impact_points:
593        for i,rays in enumerate(RayListHistory):
594            points = np.array([list(r.point) for r in rays], dtype=np.float32)
595            points = pv.PolyData(points)
596            color = pv.Color(colors[i-1])
597            fig.add_mesh(points, color=color, point_size=5, name=f"RayImpactPoints_{i}")
598    
599    detector_copies = {key: copy(OpticalChain.detectors[key]) for key in OpticalChain.detectors.keys()}
600    detector_meshes_list = []
601    detectedpoint_meshes = dict()
602    
603    if OpticalChain.detectors is not None and DrawDetectors:
604        # Detector display
605        for key in OpticalChain.detectors.keys():
606            det = detector_copies[key]
607            index = OpticalChain.detectors[key].index
608            if key in Observers:
609                det.distance = Observers[key].value
610                #Observers[key].register_calculation(lambda x: det.set_distance(x))
611            mpm._RenderDetector(fig, det, name = key, detector_meshes = detector_meshes_list)
612            if DetectedRays:
613                RayListAnalysed = OpticalChain.get_output_rays()[index]
614                points = det.get_3D_points(RayListAnalysed)
615                points = pv.PolyData(points)
616                detectedpoint_meshes[key] = points
617                fig.add_mesh(points, color='purple', point_size=5, name=f"DetectedRays_{key}")
618    detector_meshes = dict(zip(OpticalChain.detectors.keys(), detector_meshes_list))
619    
620    # Now we define a function that will move on the plot the detector with name "detname" when it's called
621    def move_detector(detname, new_value):
622        nonlocal fig, detector_meshes, detectedpoint_meshes, DetectedRays, detectedpoint_meshes, detector_copies, OpticalChain
623        det = detector_copies[detname]
624        index = OpticalChain.detectors[detname].index
625        det_mesh = detector_meshes[detname]
626        translation = det.normal * (det.distance - new_value)
627        det_mesh.translate(translation, inplace=True)
628        det.distance = new_value
629        if DetectedRays:
630            points_mesh = detectedpoint_meshes[detname]
631            points_mesh.points = det.get_3D_points(OpticalChain.get_output_rays()[index])
632        fig.show()
633    
634    # Now we register the function to the observers
635    for key in OpticalChain.detectors.keys():
636        if key in Observers:
637            Observers[key].register(lambda x: move_detector(key, x))
638
639    #pv.save_meshio('optics.inp', pointcloud)  
640    print(
641        "\r\033[K", end="", flush=True
642    )  # move to beginning of the line with \r and then delete the whole line with \033[K
643    fig.show()
644    return fig
645
646moc.OpticalChain.render = DrawSetup
647
648# %% Asphericity
649
650def DrawAsphericity(Mirror, Npoints=1000):
651    """
652    This function displays a map of the asphericity of the mirror.
653    It's a scatter plot of the points of the mirror surface, with the color representing the distance to the closest sphere.
654    The closest sphere is calculated by the function get_closest_sphere, so least square method.
655
656    Parameters
657    ----------
658    Mirror : Mirror
659        The mirror to analyse.
660
661    Npoints : int, optional
662        The number of points to sample on the mirror surface. The default is 1000.
663    
664    Returns
665    -------
666    fig : Figure
667        The figure of the plot.
668    """
669    plt.ion()
670    fig = plt.figure()
671    ax = Mirror.support._ContourSupport(fig)
672    center, radius = man.GetClosestSphere(Mirror, Npoints)
673    Points = mpm.sample_support(Mirror.support, Npoints=1000)
674    Points += Mirror.r0[:2]
675    Z = Mirror._zfunc(Points)
676    Points = mgeo.PointArray([Points[:, 0], Points[:, 1], Z]).T
677    X, Y = Points[:, 0] - Mirror.r0[0], Points[:, 1] - Mirror.r0[1]
678    Points_centered = Points - center
679    Distance = np.linalg.norm(Points_centered, axis=1) - radius
680    Distance*=1e3 # To convert to µm
681    p = plt.scatter(X, Y, c=Distance, s=15)
682    divider = man.make_axes_locatable(ax)
683    cax = divider.append_axes("right", size="5%", pad=0.05)
684    cbar = fig.colorbar(p, cax=cax)
685    cbar.set_label("Distance to closest sphere (µm)")
686    ax.set_xlabel("x (mm)")
687    ax.set_ylabel("y (mm)")
688    plt.title("Asphericity map", loc="right")
689    plt.tight_layout()
690
691    bbox = ax.get_position()
692    bbox.set_points(bbox.get_points() - np.array([[0.01, 0], [0.01, 0]]))
693    ax.set_position(bbox)
694    plt.show()
695    return fig
696
697mmirror.Mirror.drawAsphericity = DrawAsphericity
698# %% Caustics
699
700def DrawCaustics(OpticalChain, Range=1, Detector="Focus" , Npoints=1000, Nrays=1000):
701    """
702    This function displays the caustics of the rays on the detector.
703    To do so, it calculates the intersections of the rays with the detector over a 
704    range determined by the parameter Range, and then plots the standard deviation of the
705    positions in the x and y directions.
706
707    Parameters
708    ----------
709    OpticalChain : OpticalChain
710        The optical chain to analyse.
711
712    DetectorName : str
713        The name of the detector on which the caustics are calculated.
714    
715    Range : float
716        The range of the detector over which to calculate the caustics.
717
718    Npoints : int, optional
719        The number of points to sample on the detector. The default is 1000.
720    
721    Returns
722    -------
723    fig : Figure
724        The figure of the plot.
725    """
726    distances = np.linspace(-Range, Range, Npoints)
727    if isinstance(Detector, str):
728        Det = OpticalChain.detectors[Detector]
729        Index = Det.index
730    Rays = OpticalChain.get_output_rays()[Index]
731    Nrays = min(Nrays, len(Rays))
732    Rays = np.random.choice(Rays, Nrays, replace=False)
733    LocalRayList = [r.to_basis(*Det.basis) for r in Rays]
734    Points = mgeo.IntersectionRayListZPlane(LocalRayList, distances)
735    x_std = []
736    y_std = []
737    for i in range(len(distances)):
738        x_std.append(mp.StandardDeviation(Points[i][:,0]))
739        y_std.append(mp.StandardDeviation(Points[i][:,1]))
740    plt.ion()
741    fig, ax = plt.subplots()
742    ax.plot(distances, x_std, label="x std")
743    ax.plot(distances, y_std, label="y std")
744    ax.set_xlabel("Detector distance (mm)")
745    ax.set_ylabel("Standard deviation (mm)")
746    ax.legend()
747    plt.title("Caustics")
748    plt.show()
749    return fig
750
751moc.OpticalChain.drawCaustics = DrawCaustics