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