ModuleProcessing

  1"""
  2Contains processing functions used for the ray-tracing and automated (pre-)alignment of the optical chains.
  3Usually these don't need to be called by users of ART, but they may be useful.
  4
  5
  6
  7Created in Apr 2020
  8
  9@author: Anthony Guillaume + Stefan Haessler
 10"""
 11# %% Modules
 12import ARTcore.ModuleGeometry as mgeo
 13from ARTcore.ModuleGeometry import Point, Vector, Origin
 14import ARTcore.ModuleMirror as mmirror
 15import ARTcore.ModuleMask as mmask
 16import ARTcore.ModuleSource as msource
 17import ARTcore.ModuleOpticalElement as moe
 18import ARTcore.ModuleOpticalRay as mray
 19import ARTcore.ModuleSupport as msupp
 20import ARTcore.ModuleOpticalChain as moc
 21
 22import os
 23import pickle
 24
 25# import gzip
 26import lzma
 27from time import perf_counter
 28from datetime import datetime
 29from copy import copy
 30import numpy as np
 31import quaternion
 32import logging
 33
 34logger = logging.getLogger(__name__)
 35
 36
 37# %%
 38def singleOEPlacement(
 39    Optic: moe.OpticalElement,
 40    Distance: float,
 41    IncidenceAngle: float = 0,
 42    IncidencePlaneAngle: float = 0,
 43    InputRay: mray.Ray = None,
 44    AlignmentVector: str = "",
 45    PreviousIncidencePlane: mgeo.Vector = mgeo.Vector([0, 1, 0])
 46):
 47    """
 48    Automatic placement and alignment of a single optical element.
 49    The arguments are:
 50    - The optic to be placed.
 51    - The distance from the previous optic.
 52    - An incidence angle.
 53    - An incidence plane angle.
 54    - An input ray.
 55    - An alignment vector.
 56
 57    The alignment procedure is as follows:
 58    - The optic is initially placed with its center at the source position (origin point of previous master ray). 
 59    - It's oriented such that the alignment vector is antiparallel to the master ray. By default, the alignment vector is the support_normal.
 60    - The majoraxis of the optic is aligned with the incidence plane. 
 61    - The optic is rotated around the incidence plane by the incidence angle.
 62    - The optic is rotated around the master ray by the incidence plane angle.
 63    - The optic is translated along the master ray by the distance from the previous optic.
 64    - The master ray is propagated through the optic without any blocking or diffraction effects and the output ray is used as the master ray for the next optic.
 65    """
 66    # Convert angles to radian and wrap to 2pi
 67    IncidencePlaneAngle = np.deg2rad(IncidencePlaneAngle) % (2 * np.pi)
 68    IncidenceAngle = np.deg2rad(IncidenceAngle) % (2 * np.pi)
 69    assert InputRay is not None, "InputRay must be provided"
 70    OldOpticalElementCentre = InputRay.point
 71    MasterRayDirection = InputRay.vector.normalized()
 72    OpticalElementCentre = OldOpticalElementCentre + MasterRayDirection * Distance
 73
 74    logger.debug(f"Old Optical Element Centre: {OldOpticalElementCentre}")
 75    logger.debug(f"Master Ray: {InputRay}")
 76    logger.debug(f"Optical Element Centre: {OpticalElementCentre}")
 77    # for convex mirrors, rotated them by 180° while keeping same incidence plane so we reflect from the "back side"
 78    if Optic.curvature == mmirror.Curvature.CONVEX:
 79        IncidenceAngle = np.pi - IncidenceAngle
 80
 81    if hasattr(Optic, AlignmentVector):
 82        OpticVector = getattr(Optic, AlignmentVector)
 83    else:
 84        logger.warning(f"Optical Element {Optic} does not have an attribute {AlignmentVector}. Using support_normal instead.")
 85        OpticVector = Optic.support_normal_ref
 86    
 87    MajorAxis = Optic.majoraxis_ref
 88
 89    IncidencePlane = PreviousIncidencePlane.rotate(mgeo.QRotationAroundAxis(InputRay.vector, -IncidencePlaneAngle))
 90    # We calculate a quaternion that will rotate OpticVector against the master ray and MajorAxis into the incidence plane
 91    # The convention is that the MajorAxis vector points right if seen from the source with IncidencePlane pointing up
 92    # To do that, we first calculate a vector MinorAxis in the reference frame of the optic that is orthogonal to both OpticVector and MajorAxis
 93    # Then we calculate a rotation matrix that will:
 94    #   - rotate MinorAxis into the IncidencePlane direction
 95    #   - rotate OpticVector against the master ray direction
 96    #   - rotate MajorAxis into the direction of the cross product of the two previous vectors
 97    MinorAxis = -np.cross(OpticVector, MajorAxis)
 98    qIncidenceAngle = mgeo.QRotationAroundAxis(IncidencePlane, -IncidenceAngle)
 99    q = mgeo.QRotationVectorPair2VectorPair(MinorAxis, IncidencePlane, OpticVector, -MasterRayDirection)
100    q = qIncidenceAngle*q
101    Optic.q = q
102    logger.debug(f"Optic: {Optic}")
103    logger.debug(f"OpticVector: {OpticVector}")
104    logger.debug(f"Rotated OpticVector: {OpticVector.rotate(q)}")
105    logger.debug(f"MasterRayDirection: {MasterRayDirection}")
106
107    Optic.r = Origin + (OpticalElementCentre - Optic.centre - Optic.r)
108
109    NextRay = Optic.propagate_raylist(mray.RayList.from_list([InputRay]), alignment=True)[0]
110    return NextRay, IncidencePlane
111
112
113def OEPlacement(
114    OpticsList: list,
115    InitialRay:mray.Ray = None,
116    Source: msource.Source = None,
117    InputIncidencePlane: mgeo.Vector = None
118):
119    """
120    Automatic placement and alignment of the optical elements for one optical chain.
121    Returns a list of optical elements (copied)
122    """
123    if InitialRay is None:
124        if Source is None:
125            InputRay = InputRay = mray.Ray(
126            point=mgeo.Point([0, 0, 0]), 
127            vector=mgeo.Vector([1, 0, 0]),
128            path=(0.0,),
129            number=0,
130            wavelength=800e-9,
131            incidence=0.0,
132            intensity=1.0
133    )
134        else:
135            InputRay = Source.get_master_ray()
136    else:
137        InputRay = InitialRay.copy_ray()
138    if InputIncidencePlane is None:
139        InputIncidencePlane = mgeo.Vector([0, 1, 0])
140    logger.debug(f"Initial Ray: {InputRay}")
141    logger.debug(f"Initial Incidence Plane: {InputIncidencePlane}")
142    assert np.linalg.norm(np.dot(InputRay.vector, InputIncidencePlane))<1e-6, "InputIncidencePlane is not orthogonal to InputRay.vector"
143    PreviousIncidencePlane = InputIncidencePlane.copy()
144    OEList = []
145    for i in range(len(OpticsList)):
146        OEList.append(copy(OpticsList[i]["OpticalElement"]))
147        InputRay, PreviousIncidencePlane = singleOEPlacement(
148            OEList[-1],
149            OpticsList[i]["Distance"],
150            OpticsList[i]["IncidenceAngle"],
151            OpticsList[i]["IncidencePlaneAngle"],
152            InputRay,
153            OpticsList[i]["Alignment"] if "Alignment" in OpticsList[i] else "support_normal",
154            PreviousIncidencePlane
155        )
156    return OEList
157    
158
159
160# %%
161def RayTracingCalculation(
162    source_rays: mray.RayList, optical_elements: list[moe.OpticalElement], **kwargs
163) -> list[list[mray.Ray]]:
164    """
165    The actual ray-tracing calculation, starting from the list of 'source_rays',
166    and propagating them from one optical element to the next in the order of the
167    items of the list 'optical_elements'.
168
169
170    Parameters
171    ----------
172        source_rays : list[Ray-objects]
173            List of input rays.
174
175        optical_elements : list[OpticalElement-objects]
176
177
178    Returns
179    -------
180        output_rays : list[list[Ray-objects]]
181            List of lists of rays, each item corresponding to the ray-bundle *after*
182            the item with the same index in the 'optical_elements'-list.
183    """
184    output_rays = []
185
186    for k in range(0, len(optical_elements)):
187        if k == 0:
188            RayList = source_rays
189        else:
190            RayList = output_rays[k - 1]
191        RayList = optical_elements[k].propagate_raylist(RayList, **kwargs)
192        output_rays.append(RayList)
193
194    return output_rays
195
196
197# %%
198def _FindOptimalDistanceBIS(movingDetector, Amplitude, Step, RayList, OptFor, IntensityWeighted):
199    ListSizeSpot = []
200    ListDuration = []
201    ListFitness = []
202    if IntensityWeighted:
203        Weights = [k.intensity for k in RayList]
204
205    movingDetector.shiftByDistance(-Amplitude)
206    n = int(2 * Amplitude / Step)
207    for i in range(n):
208        ListPointDetector2DCentre = movingDetector.get_PointList2DCentre(RayList)
209        if OptFor in ["intensity", "spotsize"]:
210            if IntensityWeighted:
211                SpotSize = WeightedStandardDeviation(ListPointDetector2DCentre, Weights)
212            else:
213                SpotSize = StandardDeviation(ListPointDetector2DCentre)
214            ListSizeSpot.append(SpotSize)
215
216        if OptFor in ["intensity", "duration"]:
217            DelayList = movingDetector.get_Delays(RayList)
218            if IntensityWeighted:
219                Duration = WeightedStandardDeviation(DelayList, Weights)
220            else:
221                Duration = StandardDeviation(DelayList)
222            ListDuration.append(Duration)
223
224        if OptFor == "intensity":
225            Fitness = SpotSize**2 * Duration
226        elif OptFor == "duration":
227            Fitness = Duration
228        elif OptFor == "spotsize":
229            Fitness = SpotSize
230        ListFitness.append(Fitness)
231
232        movingDetector.shiftByDistance(Step)
233
234    FitnessMin = min(ListFitness)
235    ind = ListFitness.index(FitnessMin)
236    if OptFor in ["intensity", "spotsize"]:
237        OptSizeSpot = ListSizeSpot[ind]
238    else:
239        OptSizeSpot = np.nan
240    if OptFor in ["intensity", "duration"]:
241        OptDuration = ListDuration[ind]
242    else:
243        OptDuration = np.nan
244
245    movingDetector.shiftByDistance(-(n - ind) * Step)
246
247    return movingDetector, OptSizeSpot, OptDuration
248
249
250def FindOptimalDistance(
251    Detector,
252    RayList: list[mray.Ray],
253    OptFor="intensity",
254    Amplitude: float = None,
255    Precision: int = 3,
256    IntensityWeighted=False,
257    verbose=False,
258):
259    """
260    Automatically finds the optimal the 'Detector'-distance for either
261    maximum intensity, or minimal spotsize or duration.
262
263    Parameters
264    ----------
265        Detector : Detector-object
266
267        RayList : list[Ray-objects]
268
269        OptFor : str, optional
270            "spotsize": minimizes the standard-deviation spot size *d* on the detector.
271            "duration": minimizes the standard-deviation *tau* of the ray-delays.
272            "intensity": Maximizes 1/tau/d^2.
273            Defaults to "intensity".
274
275        Amplitude : float, optional
276            The detector-distances within which the optimization is done will be
277            the distance of 'Detector' +/- Amplitude in mm.
278
279        Precision : int, optional
280            Precision-parameter for the search algorithm. For Precision = n,
281            it will iterate with search amplitudes decreasing until 10^{-n}.
282            Defaults to 3.
283
284        IntensityWeighted : bool, optional
285            Whether to weigh the calculation of spotsize and/or duration by the ray-intensities.
286            Defaults to False.
287
288        verbose : bool, optional
289            Whether to print results to the console.
290
291    Returns
292    -------
293        OptDetector : Detector-object
294            A new detector-instance with optimized distance.
295
296        OptSpotSize : float
297            Spotsize, calculated as standard deviation in mm of the spot-diagram
298            on the optimized detector.
299
300        OptDuration : float
301            Duration, calculated as standard deviation in fs of the ray-delays
302            on the optimized detector.
303    """
304
305    if OptFor not in ["intensity", "size", "duration"]:
306        raise NameError(
307            "I don`t recognize what you want to optimize the detector distance for. OptFor must be either 'intensity', 'size' or 'duration'."
308        )
309
310    FirstDistance = Detector.get_distance()
311    ListPointDetector2DCentre = Detector.get_PointList2DCentre(RayList)
312    SizeSpot = 2 * StandardDeviation(ListPointDetector2DCentre)
313    NumericalAperture = ReturnNumericalAperture(RayList, 1)
314    if Amplitude is None:
315        Amplitude = min(4 * np.ceil(SizeSpot / np.tan(np.arcsin(NumericalAperture))), FirstDistance)
316    Step = Amplitude/10
317    #Step = Amplitude / 5  # pretty good in half the time ;-)
318
319    if verbose:
320        print(
321            f"Searching optimal detector position for *{OptFor}* within [{FirstDistance-Amplitude:.3f}, {FirstDistance+Amplitude:.3f}] mm...",
322            end="",
323            flush=True,
324        )
325    movingDetector = Detector.copy_detector()
326    for k in range(Precision + 1):
327        movingDetector, OptSpotSize, OptDuration = _FindOptimalDistanceBIS(
328            movingDetector, Amplitude * 0.1**k, Step * 0.1**k, RayList, OptFor, IntensityWeighted
329        )
330
331    if (
332        not FirstDistance - Amplitude + 10**-Precision
333        < movingDetector.get_distance()
334        < FirstDistance + Amplitude - 10**-Precision
335    ):
336        print("There`s no minimum-size/duration focus in the searched range.")
337
338    print(
339        "\r\033[K", end="", flush=True
340    )  # move to beginning of the line with \r and then delete the whole line with \033[K
341    return movingDetector, OptSpotSize, OptDuration
342
343
344# %%
345def FindCentralRay(RayList: list[mray.Ray]):
346    """
347    Out of a the list of Ray-objects, RayList, determine the average direction
348    vector and source point, and the return a Ray-object representing this 
349    central ray of the RayList.
350    
351    Parameters
352    ----------
353        RayList : list of Ray-objects
354
355    Returns
356    -------
357        CentralRay : Ray-object
358    """
359      
360    CentralVector = np.mean( [x.vector for x in RayList], axis=0)
361    CentralPoint = np.mean( [x.point for x in RayList], axis=0)
362    
363    return mray.Ray(mgeo.Point(CentralPoint), mgeo.Vector(CentralVector))
364
365
366def StandardDeviation(List: list[float, np.ndarray]) -> float:
367    """
368    Return the standard deviation of elements in the supplied list.
369    If the elements are floats, it's simply the root-mean-square over the numbers.
370    If the elements are vectors (represented as numpy-arrays), the root-mean-square
371    of the distance of the points from their mean is calculated.
372
373    Parameters
374    ----------
375        List : list of floats or of np.ndarray
376            Numbers (in ART, typically delays)
377            or 2D or 3D vectors (in ART typically space coordinates)
378
379    Returns
380    -------
381        Std : float
382    """
383    if type(List[0]) in [int, float, np.float64]:
384        return np.std(List)
385    elif len(List[0]) > 1:
386        return np.sqrt(np.var(List, axis=0).sum())
387    else:
388        raise ValueError("StandardDeviation expects a list of floats or numpy-arrays as input, but got something else.")
389
390
391def WeightedStandardDeviation(List: list[float, np.ndarray], Weights: list[float]) -> float:
392    """
393    Return the weighted standard deviation of elements in the supplied list.
394    If the elements are floats, it's simply the root-mean-square over the weighted numbers.
395    If the elements are vectors (represented as numpy-arrays), the root-mean-square
396    of the weighted distances of the points from their mean is calculated.
397
398    Parameters
399    ----------
400        List : list of floats or of np.ndarray
401            Numbers (in ART, typically delays)
402            or 2D or 3D vectors (in ART typically space coordinates)
403
404        Weights : list of floats, same length as List
405            Parameters such as Ray.intensity to be used as weights
406
407    Returns
408    -------
409        Std : float
410    """
411    average = np.average(List, axis=0, weights=Weights, returned=False)
412    variance = np.average((List - average) ** 2, axis=0, weights=Weights, returned=False)
413    return np.sqrt(variance.sum())
414
415
416# %%
417def _hash_list_of_objects(list):
418    """returns a hash value for a list of objects, by just summing up all the individual hashes."""
419    total_hash = 0
420    for x in list:
421        total_hash += hash(x)
422    return total_hash
423
424
425def _which_indeces(lst):
426    """takes an input-list, and gives you a list of the indeces where the input-list contains another list or a numpy-array"""
427    indexes = [i for i, x in enumerate(lst) if isinstance(x, (list, np.ndarray))]
428    return indexes
429
430
431# %%
432def save_compressed(obj, filename: str = None):
433    """Save (=pickle) an object 'obj' to a compressed file with name 'filename'."""
434    if not type(filename) == str:
435        filename = "kept_data_" + datetime.now().strftime("%Y-%m-%d-%Hh%M")
436
437    i = 0
438    while os.path.exists(filename + f"_{i}.xz"):
439        i += 1
440    filename = filename + f"_{i}"
441    # with gzip.open(filename + '.gz', 'wb') as f:
442    with lzma.open(filename + ".xz", "wb") as f:
443        pickle.dump(obj, f)
444    print("Saved results to " + filename + ".xz.")
445    print("->To reload from disk do: kept_data = mp.load_compressed('" + filename + "')")
446
447
448def load_compressed(filename: str):
449    """Load (=unpickle) an object 'obj' from a compressed file with name 'filename'."""
450    # with gzip.open(filename + '.gz', 'rb') as f:
451    with lzma.open(filename + ".xz", "rb") as f:
452        obj = pickle.load(f)
453    return obj
454
455
456# %%  tic-toc timer functions
457_tstart_stack = []
458
459
460def _tic():
461    _tstart_stack.append(perf_counter())
462
463
464def _toc(fmt="Elapsed: %s s"):
465    print(fmt % (perf_counter() - _tstart_stack.pop()))