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 wavelength.
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            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