1"""
  2Adds various analysis methods.
  3
  4Contains the following functions:
  5    - GetETransmission: Calculates the energy transmission from RayListIn to RayListOut in percent.
  6    - GetResultSummary: Calculate and return FocalSpotSize-standard-deviation and Duration-standard-deviation.
  7    - GetPulseProfile: Retrieve the pulse profile from the given RayList and Detector.
  8    - GetNumericalAperture: Returns the numerical aperture associated with the supplied ray-bundle.
  9    - GetAiryRadius: Returns the radius of the Airy disk.
 10    - GetPlaneWaveFocus: Calculates the approximate polychromatic focal spot of a set of rays.
 11    - GetDiffractionFocus: Calculates the approximate polychromatic focal spot of a set of rays.
 12    - GetClosestSphere: Calculates the closest sphere to the surface of a mirror.
 13    - GetAsphericity: Calculates the maximum distance of the mirror surface to the closest sphere.
 14    - _best_fit_sphere: Calculates the best sphere to fit a set of points.
 15
 16Adds the following methods:
 17    - OpticalChain:
 18        - getETransmission: Calculates the energy transmission from the input to the output of the OpticalChain.
 19        - getResultsSummary: Calculate and return FocalSpotSize-standard-deviation and Duration-standard-deviation.
 20        - getPulseProfile: Retrieve the pulse profile from the output of the OpticalChain.
 21        - getPlaneWaveFocus: Calculates the approximate polychromatic focal spot of the output of the OpticalChain.
 22        - getDiffractionFocus: Calculates the approximate polychromatic focal spot of the output of the OpticalChain.
 23    - Mirror:
 24        - getClosestSphere: Calculates the closest sphere to the surface of a mirror.
 25        - getAsphericity: Calculates the maximum distance of the mirror surface to the closest sphere.
 26
 27Created in July 2024
 28
 29@author: André Kalouguine + Stefan Haessler + Anthony Guillaume
 30"""
 31# %% Module imports
 32import ARTcore.ModuleGeometry as mgeo
 33import ARTcore.ModuleOpticalRay as mray
 34import ARTcore.ModuleProcessing as mp
 35import ARTcore.ModuleOpticalChain as moc
 36import ARTcore.ModuleMirror as mmirror
 37import ART.ModulePlottingMethods as mpm
 38import matplotlib.pyplot as plt
 39from mpl_toolkits.axes_grid1 import make_axes_locatable
 40import numpy as np
 41import math
 42
 43LightSpeed = 299792458000
 44
 45# %% Analysis methods
 46def GetETransmission(RayListIn, RayListOut) -> float:
 47    """
 48    Calculates the energy transmission from RayListIn to RayListOut in percent by summing up the
 49    intensity-property of the individual rays.
 50
 51    Parameters
 52    ----------
 53        RayListIn : list(Ray)
 54            List of incoming rays.
 55        
 56        RayListOut : list(Ray)
 57            List of outgoing rays.
 58
 59    Returns
 60    -------
 61        ETransmission : float
 62    """
 63    ETransmission = 100 * sum(Ray.intensity for Ray in RayListOut) / sum(Ray.intensity for Ray in RayListIn)
 64    return ETransmission
 65
 66def _getETransmission(OpticalChain, IndexIn=0, IndexOut=-1) -> float:
 67    """
 68    Calculates the energy transmission from the input to the output of the OpticalChain in percent.
 69
 70    Parameters
 71    ----------
 72        OpticalChain : OpticalChain
 73            An object of the ModuleOpticalChain.OpticalChain-class.
 74
 75        IndexIn : int, optional
 76            Index of the input RayList in the OpticalChain, defaults to 0.
 77
 78        IndexOut : int, optional
 79            Index of the output RayList in the OpticalChain, defaults to -1.
 80
 81    Returns
 82    -------
 83        ETransmission : float
 84    """
 85    Rays = OpticalChain.get_output_rays()
 86    if IndexIn == 0:
 87        RayListIn = OpticalChain.get_input_rays()
 88    else:
 89        RayListIn = Rays[IndexIn]
 90    RayListOut = Rays[IndexOut]
 91    ETransmission = GetETransmission(RayListIn, RayListOut)
 92    return ETransmission
 93
 94moc.OpticalChain.getETransmission = _getETransmission
 95
 96def GetResultSummary(Detector, RayListAnalysed, verbose=False):
 97    """
 98    Calculate and return FocalSpotSize-standard-deviation and Duration-standard-deviation
 99    for the given Detector and RayList.
100    If verbose, then also print a summary of the results for the given Detector.
101
102    Parameters
103    ----------
104        Detector : Detector
105            An object of the ModuleDetector.Detector-class.
106
107        RayListAnalysed : list(Ray)
108            List of objects of the ModuleOpticalRay.Ray-class.
109
110        verbose : bool
111            Whether to print a result summary.
112
113    Returns
114    -------
115        FocalSpotSizeSD : float
116
117        DurationSD : float
118    """
119    DetectorPointList2DCentre = Detector.get_2D_points(RayListAnalysed)
120    FocalSpotSizeSD = mp.StandardDeviation(DetectorPointList2DCentre)
121    DelayList = Detector.get_Delays(RayListAnalysed)
122    DurationSD = mp.StandardDeviation(DelayList)
123
124    if verbose:
125        FocalSpotSize = mgeo.DiameterPointList(DetectorPointList2DCentre)
126        summarystring = (
127            "At the detector distance of "
128            + "{:.3f}".format(Detector.get_distance())
129            + " mm we get:\n"
130            + "Spatial std : "
131            + "{:.3f}".format(FocalSpotSizeSD * 1e3)
132            + " \u03BCm and min-max: "
133            + "{:.3f}".format(FocalSpotSize * 1e3)
134            + " \u03BCm\n"
135            + "Temporal std : "
136            + "{:.3e}".format(DurationSD)
137            + " fs and min-max : "
138            + "{:.3e}".format(max(DelayList) - min(DelayList))
139            + " fs"
140        )
141
142        print(summarystring)
143
144    return FocalSpotSizeSD, DurationSD
145
146def _getResultsSummary(OpticalChain, Detector = "Focus", verbose=False):
147    """
148    Calculate and return FocalSpotSize-standard-deviation and Duration-standard-deviation
149    for the given Detector and RayList.
150    If verbose, then also print a summary of the results for the given Detector.
151
152    Parameters
153    ----------
154        OpticalChain : OpticalChain
155            An object of the ModuleOpticalChain.OpticalChain-class.
156
157        Detector : Detector or str, optional
158            An object of the ModuleDetector.Detector-class or "Focus" to use the focus detector, defaults to "Focus".
159
160        verbose : bool
161            Whether to print a result summary.
162
163    Returns
164    -------
165        FocalSpotSizeSD : float
166
167        DurationSD : float
168    """
169    if isinstance(Detector, str):
170        Detector = OpticalChain.detectors[Detector]
171    Index = Detector.index
172    RayListAnalysed = OpticalChain.get_output_rays()[Index]
173    FocalSpotSizeSD, DurationSD = GetResultSummary(Detector, RayListAnalysed, verbose)
174    return FocalSpotSizeSD, DurationSD
175
176moc.OpticalChain.getResultsSummary = _getResultsSummary
177
178def GetPulseProfile(Detector, RayList, Nbins=100):
179    """
180    Retrieve the pulse profile from the given RayList and Detector.
181    The pulse profile is calculated by binning the delays of the rays in the RayList
182    at the Detector. The intensity of the rays is also taken into account.
183
184    Parameters
185    ----------
186        Detector : Detector
187            An object of the ModuleDetector.Detector-class.
188
189        RayList : list(Ray)
190            List of objects of the ModuleOpticalRay.Ray-class.
191
192        Nbins : int
193            Number of bins to use for the histogram.
194
195    Returns
196    -------
197        time : numpy.ndarray
198            The bins of the histogram.
199
200        intensity : numpy.ndarray
201            The histogram values.
202    """
203    DelayList = Detector.get_Delays(RayList)
204    IntensityList = [Ray.intensity for Ray in RayList]
205    intensity, time_edges = np.histogram(DelayList, bins=Nbins, weights=IntensityList)
206    time = 0.5 * (time_edges[1:] + time_edges[:-1])
207    return time, intensity
208
209def _getPulseProfile(OpticalChain, Detector = "Focus", Nbins=100):
210    """
211    Retrieve the pulse profile from the output of the OpticalChain.
212    The pulse profile is calculated by binning the delays of the rays in the RayList
213    at the Detector. The intensity of the rays is also taken into account.
214
215    Parameters
216    ----------
217        OpticalChain : OpticalChain
218            An object of the ModuleOpticalChain.OpticalChain-class.
219
220        Detector : Detector or str, optional
221            An object of the ModuleDetector.Detector-class or "Focus" to use the focus detector, defaults to "Focus".
222
223        Nbins : int
224            Number of bins to use for the histogram.
225
226    Returns
227    -------
228        time : numpy.ndarray
229            The bins of the histogram.
230
231        intensity : numpy.ndarray
232            The histogram values.
233    """
234    if isinstance(Detector, str):
235        Detector = OpticalChain.detectors[Detector]
236    Index = Detector.index
237    RayList = OpticalChain.get_output_rays()[Index]
238    time, intensity = GetPulseProfile(Detector, RayList, Nbins)
239    return time, intensity
240
241
242def GetNumericalAperture(RayList: list[mray.Ray], RefractiveIndex: float = 1) -> float:
243    r"""
244    Returns the numerical aperture associated with the supplied ray-bundle 'Raylist'.
245    This is $n\sin\theta$, where $\theta$ is the maximum angle between any of the rays and the central ray,
246    and $n$ is the refractive index of the propagation medium.
247
248    Parameters
249    ----------
250        RayList : list of Ray-object
251            The ray-bundle of which to determine the NA.
252
253        RefractiveIndex : float, optional
254            Refractive index of the propagation medium, defaults to =1.
255
256    Returns
257    -------
258        NA : float
259    """
260    CentralRay = mp.FindCentralRay(RayList)
261    if CentralRay is None:
262        CentralVector = np.array([0, 0, 0])
263        for k in RayList:
264            CentralVector = CentralVector + k.vector
265        CentralVector = CentralVector / len(RayList)
266    else:
267        CentralVector = CentralRay.vector
268    ListAngleAperture = []
269    for k in RayList:
270        ListAngleAperture.append(mgeo.AngleBetweenTwoVectors(CentralVector, k.vector))
271
272    return np.sin(np.amax(ListAngleAperture)) * RefractiveIndex
273
274
275def GetAiryRadius(Wavelength: float, NumericalAperture: float) -> float:
276    r"""
277    Returns the radius of the Airy disk: $r = 1.22 \frac\{\lambda\}\{NA\}$,
278    i.e. the diffraction-limited radius of the focal spot corresponding to a given
279    numerical aperture $NA$ and a light wavelength $\lambda$.
280
281    For very small $NA<10^\{-3\}$, diffraction effects becomes negligible and the Airy Radius becomes meaningless,
282    so in that case, a radius of 0 is returned.
283
284    Parameters
285    ----------
286        Wavelength : float
287            Light wavelength in mm.
288
289        NumericalAperture : float
290
291    Returns
292    -------
293        AiryRadius : float
294    """
295    if NumericalAperture > 1e-3 and Wavelength is not None:
296        return 1.22 * 0.5 * Wavelength / NumericalAperture
297    else:
298        return 0  # for very small numerical apertures, diffraction effects becomes negligible and the Airy Radius becomes meaningless
299
300
301def GetPlaneWaveFocus(OpticalChain, Detector = "Focus", size=None, Nrays=1000, resolution=100):
302    """
303    This function calculates the approximate polychromatic focal spot of a set of rays.
304    To do so, it positions itself in the detector plane and samples a square area.
305    If the size of the area is not given, it will use twice the Airy radius of the system with the largest wavbelength.
306    Otherwise it uses the size.
307    The resolution of the sampling is given by the resolution parameter. So it samples on a grid resolution x resolution.
308
309    To calculate the intensity, it takes Nrays out of the raylist (to subsample if needed).
310    It assimilates every ray to a plane wave, so it calculates a k-vector for each ray (taking into account the wavelength of the ray).
311    It calculates the phase of each ray from the delay and from the intersection position with the detector.
312    The delay from the non-central position is simply sin(alpha)*distance/c, where alpha is the angle between the ray and the normal to the detector
313    and distance is the distance between the intersection point and the central point of the detector. 
314    It then calculates the intensity at each point of the grid by summing the intensity of each plane wave.
315
316    As long as there are not too many rays, this method is faster than doing an FFT.
317    It's also naturally polychromatic.
318    """
319    if isinstance(Detector, str):
320        Detector = OpticalChain.detectors[Detector]
321    Index = Detector.index
322    RayList = OpticalChain.get_output_rays()[Index]
323    if size is None:
324        Wavelengths = [Ray.wavelength for Ray in RayList]
325        Wavelength = max(Wavelengths)
326        NumericalAperture = GetNumericalAperture(RayList)
327        size = 3 * GetAiryRadius(Wavelength, NumericalAperture)
328    X = np.linspace(-size / 2, size / 2, resolution)
329    Y = np.linspace(-size / 2, size / 2, resolution)
330    X, Y = np.meshgrid(X, Y)
331    # We pick Nrays. if there are less than Nrays, we take all of them.
332    PickedRaysGlobal = np.random.choice(RayList, min(Nrays, len(RayList)), replace=False)
333    # We calculate the k-vector of each ray, taking into account the wavelength of the ray.
334    # The units should be in mm^-1
335    PickedRays = [Ray.to_basis(*Detector.basis) for Ray in PickedRaysGlobal]
336    # The rays are now in the reference plane of the detector whose normal is [0,0,1]
337    wavelengths = np.array([Ray.wavelength for Ray in PickedRays])
338    vectors = np.array([Ray.vector for Ray in PickedRays])
339    frequencies = np.array([LightSpeed / Ray.wavelength for Ray in PickedRays]) # mm/s / mm = 1/s
340    k_vectors = 2*np.pi * vectors / wavelengths[:, np.newaxis]
341    angles = np.arccos(np.clip(np.dot(vectors, np.array([0, 0, 1])), -1.0, 1.0))
342    # We calculate the intersection of the rays with the detector plane
343    Intersections = mgeo.IntersectionRayListZPlane(PickedRays)[0]
344    distances = (Intersections - mgeo.Origin[:2]).norm
345    # We calculate the phase of each ray
346    PathDelays = np.array(Detector.get_Delays(PickedRaysGlobal))
347    PositionDelays = np.sum(np.array(Intersections._add_dimension()-mgeo.Origin)*vectors, axis=1)  / LightSpeed * 1e15
348    Delays = (PathDelays+PositionDelays) /  1e15
349    # We have the delays in fs, we can now calculate the phase, taking into account the wavelength of each ray
350    Phases = np.array([2 * np.pi * frequencies[i] * Delays[i] for i in range(len(Delays))])
351    # We also need the intensity of each ray
352    RayIntensities = np.array([Ray.intensity for Ray in PickedRays])
353    # We can now calculate the intensity at each point of the grid
354    Intensity = np.zeros((resolution, resolution))
355    for i in range(len(PickedRays)):
356        Intensity += RayIntensities[i] * np.cos(k_vectors[i][0] * X + k_vectors[i][1] * Y + Phases[i])
357
358    return X, Y, Intensity
359
360moc.getPlaneWaveFocus = GetPlaneWaveFocus
361
362def GetDiffractionFocus(OpticalChain, Detector = "Focus", size=None, Nrays=1000, resolution=100):
363    """
364    This function calculates the approximate polychromatic focal spot of a set of rays.
365    To do so, it positions itself in the detector plane and samples a square area.
366    If the size of the area is not given, it will use twice the Airy radius of the system with the largest wavbelength.
367    Otherwise it uses the size.
368    The resolution of the sampling is given by the resolution parameter. So it samples on a grid resolution x resolution.
369
370    To calculate the intensity, it takes Nrays out of the raylist (to subsample if needed).
371    It assimilates every ray to a plane wave, so it calculates a k-vector for each ray (taking into account the wavelength of the ray).
372    It considers that all the rays are intersecting the detector in the middle and doesn't take into account their phase.
373
374    So it returns the best case scenario for a diffraction limited focus with that numerical aperture
375    """
376    if isinstance(Detector, str):
377        Detector = OpticalChain.detectors[Detector]
378    Index = Detector.index
379    RayList = OpticalChain.get_output_rays()[Index]
380    if size is None:
381        Wavelengths = [Ray.wavelength for Ray in RayList]
382        Wavelength = max(Wavelengths)
383        NumericalAperture = GetNumericalAperture(RayList)
384        size = 3 * GetAiryRadius(Wavelength, NumericalAperture)
385    X = np.linspace(-size / 2, size / 2, resolution)
386    Y = np.linspace(-size / 2, size / 2, resolution)
387    X, Y = np.meshgrid(X, Y)
388    # We pick Nrays. if there are less than Nrays, we take all of them.
389    PickedRaysGlobal = np.random.choice(RayList, min(Nrays, len(RayList)), replace=False)
390    # We calculate the k-vector of each ray, taking into account the wavelength of the ray.
391    # The units should be in mm^-1
392    PickedRays = [Ray.to_basis(*Detector.basis) for Ray in PickedRaysGlobal]
393    # The rays are now in the reference plane of the detector whose normal is [0,0,1]
394    wavelengths = np.array([Ray.wavelength for Ray in PickedRays])
395    vectors = np.array([Ray.vector for Ray in PickedRays])
396    frequencies = np.array([LightSpeed / Ray.wavelength for Ray in PickedRays]) # mm/s / mm = 1/s
397    k_vectors = 2*np.pi * vectors / wavelengths[:, np.newaxis]
398    angles = np.arccos(np.clip(np.dot(vectors, np.array([0, 0, 1])), -1.0, 1.0))
399    # We calculate the intersection of the rays with the detector plane
400    Intersections = mgeo.IntersectionRayListZPlane(PickedRays)[0]
401    distances = (Intersections - mgeo.Origin[:2]).norm
402    # We also need the intensity of each ray
403    RayIntensities = np.array([Ray.intensity for Ray in PickedRays])
404    # We can now calculate the intensity at each point of the grid
405    Intensity = np.zeros((resolution, resolution))
406    for i in range(len(PickedRays)):
407        Intensity += RayIntensities[i] * np.cos(k_vectors[i][0] * X + k_vectors[i][1] * Y)
408
409    return X, Y, Intensity
410
411moc.getDiffractionFocus = GetDiffractionFocus
412
413# %% Asphericity analysis
414# This code is to do asphericity analysis of various surfaces
415# It calculates the closes sphere to the surface of an optical element
416# Then there are two functions. One that simply gives an asphericity value
417# and another one that actually plots the distance the the closest sphere 
418# in much the same way as we plot the MirrorProjection
419
420def BestFitSphere(X,Y,Z):
421    """
422    This function calculates the best sphere to fit a set of points.
423    It uses the least square method to find the center and the radius of the sphere.
424    Cite https://jekel.me/2015/Least-Squares-Sphere-Fit/
425    """
426    A = np.zeros((len(X),4))
427    A[:,0] = X*2
428    A[:,1] = Y*2
429    A[:,2] = Z*2
430    A[:,3] = 1
431
432    #   Assemble the f matrix
433    f = np.zeros((len(X),1))
434    f[:,0] = X**2 + Y**2 + Z**2
435    C, residules, rank, singval = np.linalg.lstsq(A,f)
436
437    #   solve for the radius
438    t = (C[0]*C[0])+(C[1]*C[1])+(C[2]*C[2])+C[3]
439    radius = math.sqrt(t)
440
441    return mgeo.Point(C[:3].flatten()),radius
442
443
444def GetClosestSphere(Mirror, Npoints=1000):
445    """
446    This function calculates the closest sphere to the surface of a mirror.
447    It does so by sampling the surface of the mirror at Npoints points.
448    It then calculates the closest sphere to these points.
449    It returns the radius of the sphere and the center of the sphere.
450    """
451    Points = mpm.sample_support(Mirror.support, Npoints=1000)
452    Points += Mirror.r0[:2]
453    Z = Mirror._zfunc(Points)
454    Points = mgeo.PointArray([Points[:, 0], Points[:, 1], Z]).T
455    spX, spY, spZ = Points[:, 0], Points[:, 1], Points[:, 2]
456    Center, Radius = BestFitSphere(spX, spY, spZ)
457    return Center, Radius
458
459def GetAsphericity(Mirror, Npoints=1000):
460    """
461    This function calculates the maximum distance of the mirror surface to the closest sphere. 
462    """
463    center, radius = GetClosestSphere(Mirror, Npoints)
464    Points = mpm.sample_support(Mirror.support, Npoints=1000)
465    Points += Mirror.r0[:2]
466    Z = Mirror._zfunc(Points)
467    Points = mgeo.PointArray([Points[:, 0], Points[:, 1], Z]).T
468    Points_centered = Points - center
469    Distance = np.linalg.norm(Points_centered, axis=1) - radius
470    Distance*=1e3 # To convert to µm
471    return np.ptp(Distance)
472
473mmirror.Mirror.getClosestSphere = GetClosestSphere
474mmirror.Mirror.getAsphericity = GetAsphericity

  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