ModuleProcessing

134def OEPlacement(
135    SourceProperties: dict,
136    OpticsList: list,
137    DistanceList: list[(int, float)],
138    IncidenceAngleList: list[float],
139    IncidencePlaneAngleList: list[float] = None,
140    Description: str = "",
141    render: bool = False,
142):
143    """
144    Automatically place optical elements in the "lab frame" according to given distances and incidence angles.
145    Outputs an OpticalChain-object.
146
147    The source is placed at the origin, and points into the x-direction. The optical elements then follow.
148
149    As long as the angles in IncidencePlaneAngleList are 0, the incidence plane remains the x-z plane, i.e.
150    the optical elements are rotated about the y-axis to set the desired incidence angle between OE-normal
151    and the direction of incidence of the incoming ray-bundle. Otherwise, the incidence plane gets rotated.
152
153    One of the elements of one of the lists 'DistanceList', 'IncidenceAngleList', or 'IncidencePlaneAngleList'
154    can be a list or a numpy-array. In that case, a list of OpticalChain-objects is created which can be looped over.
155
156    Parameters
157    ----------
158        SourceProperties : dict
159            Dictionary with the keys "Divergence" :float, "SourceSize" :float,
160            "NumberRays" :int, and "Wavelength" :float.
161            This determines the source ray-bundle that is created.
162
163        OpticsList : list[Optics-class-objects (Mirorrs or Masks)]
164            List of Optics-objects.
165
166        DistanceList : list[float]
167            List of distances, in mm,  at which place the optics out of the
168            OpticsList from the precending one. For the first optic, this is
169            the distance from the light source.
170
171            One of the elements can ne a list or a numpy-array instead of a float.
172            In that case, a list of OpticalChain-objects is created which can be looped over.
173
174        IncidenceAngleList : list[float]
175            List of incidence angles, in degrees, measured between the central ray
176            of the incident ray bundle and the normal on the optical element.
177
178            One of the elements can ne a list or a numpy-array instead of a float.
179            In that case, a list of OpticalChain-objects is created which can be looped over.
180
181        IncidencePlaneAngleList : list[float], optional
182            List of angles, in degrees, by which the incidence plane is rotated
183            on the optical element with respect to that on the preceding one.
184            Consequently, for the first element, this angle has no qualitative effect.
185
186            One of the elements can ne a list or a numpy-array instead of a float.
187            In that case, a list of OpticalChain-objects is created which can be looped over.
188
189        Description : str, optional
190            A string to describe the optical setup.
191
192        render : bool, optional
193            Whether save a png-image with a rendering of the optical setup.
194
195    Returns
196    -------
197        OpticalChainList : list[OpticalChain-object]
198
199        or
200
201        OpticalChain : OpticalChain-object
202
203    """
204    if render:
205        from mayavi import mlab
206
207    if IncidencePlaneAngleList is None:
208        IncidencePlaneAngleList = np.zeros(len(OpticsList))
209
210    nest_indx_distance = _which_indeces(DistanceList)
211    nest_indx_incidence = _which_indeces(IncidenceAngleList)
212    nest_indx_incplane = _which_indeces(IncidencePlaneAngleList)
213
214    total_nested = len(nest_indx_incidence + nest_indx_incplane + nest_indx_distance)
215    if total_nested > 1:
216        raise ValueError(
217            "Only one element of one of the lists IncidenceAngleList, IncidencePlaneAngleList, or DistanceList can be a list or array itself. Otherwise things get too tangled..."
218        )
219    elif total_nested == 1:
220        i = (nest_indx_incidence + nest_indx_incplane + nest_indx_distance)[0]
221        loop_variable_name = OpticsList[i].type + "_idx_" + str(i)
222        if nest_indx_incidence:  # if the list is not empty
223            loop_variable_name += " incidence angle (deg)"
224            loop_variable = copy.deepcopy(IncidenceAngleList[i])
225            loop_list = IncidenceAngleList
226        elif nest_indx_distance:  # if the list is not empty
227            loop_variable_name += " distance (mm)"
228            loop_variable = copy.deepcopy(DistanceList[i])
229            loop_list = DistanceList
230        elif nest_indx_incplane:  # if the list is not empty
231            loop_variable_name += " incidence-plane angle rotation (deg)"
232            loop_variable = copy.deepcopy(IncidencePlaneAngleList[i])
233            loop_list = IncidencePlaneAngleList
234
235        OpticalChainList = []
236        for x in loop_variable:
237            loop_list[i] = x
238            ModifiedOpticalChain = _singleOEPlacement(
239                SourceProperties, OpticsList, DistanceList, IncidenceAngleList, IncidencePlaneAngleList, Description
240            )
241            ModifiedOpticalChain.loop_variable_name = loop_variable_name
242            ModifiedOpticalChain.loop_variable_value = x
243            OpticalChainList.append(ModifiedOpticalChain)
244
245            if render:
246                # to produce and save renderings of the setup for each iteration in the loop:
247                mlab.options.offscreen = True
248                fig = ModifiedOpticalChain.render()
249                name = loop_variable_name + "=" + "{:.2f}".format(x) + ".png"
250                name = name.replace(" ", "_")
251                print("...saving image...", end="", flush=True)
252                mlab.savefig(name, figure=fig, magnification=3)
253                print(
254                    "\r\033[K", end="", flush=True
255                )  # move to beginning of the line with \r and then delete the whole line with \033[K
256
257        if render:
258            mlab.options.offscreen = False
259
260        return OpticalChainList
261
262    elif total_nested == 0:
263        OpticalChain = _singleOEPlacement(
264            SourceProperties, OpticsList, DistanceList, IncidenceAngleList, IncidencePlaneAngleList, Description
265        )  # all simple
266        if render:
267            mlab.options.offscreen = True
268            fig = OpticalChain.render()
269            name = "OpticalChain_" + datetime.now().strftime("%Y-%m-%d-%Hh%M") + ".png"
270            mlab.savefig(name, figure=fig, magnification=3)
271            mlab.options.offscreen = False
272
273        return OpticalChain

277def RayTracingCalculation(
278    source_rays: list[mray.Ray], optical_elements: list[moe.OpticalElement], IgnoreDefects = True
279) -> list[list[mray.Ray]]:
280    """
281    The actual ray-tracing calculation, starting from the list of 'source_rays',
282    and propagating them from one optical element to the next in the order of the
283    items of the list 'optical_elements'.
284
285
286    Parameters
287    ----------
288        source_rays : list[Ray-objects]
289            List of input rays.
290
291        optical_elements : list[OpticalElement-objects]
292
293
294    Returns
295    -------
296        output_rays : list[list[Ray-objects]]
297            List of lists of rays, each item corresponding to the ray-bundle *after*
298            the item with the same index in the 'optical_elements'-list.
299    """
300    ez = np.array([0, 0, 1])
301    ex = np.array([1, 0, 0])
302    output_rays = []
303
304    for k in range(0, len(optical_elements)):
305        if k == 0:
306            RayList = source_rays
307        else:
308            RayList = output_rays[k - 1]
309
310        # Mirror = optical_elements[k].type
311        Position = optical_elements[k].position
312        n = optical_elements[k].normal
313        m = optical_elements[k].majoraxis
314
315        # transform rays into mirror coordinate system
316        RayList = mgeo.TranslationRayList(RayList, -Position)
317        RayList = mgeo.RotationRayList(RayList, n, ez)
318        # FIRST rotate m in the same way as you just did the ray list,
319        # THEN rotate the rays again to match the NEW m with the x-axis !
320        mPrime = mgeo.RotationPoint(m, n, ez)
321        RayList = mgeo.RotationRayList(RayList, mPrime, ex)
322        RayList = mgeo.TranslationRayList(RayList, optical_elements[k].type.get_centre())
323
324        # optical element acts on the rays:
325        if "Mirror" in optical_elements[k].type.type:
326            RayList = mmirror.ReflectionMirrorRayList(optical_elements[k].type, RayList, IgnoreDefects = IgnoreDefects)
327        elif optical_elements[k].type.type == "Mask":
328            RayList = mmask.TransmitMaskRayList(optical_elements[k].type, RayList)
329        else:
330            raise NameError("I don`t recognize the type of optical element " + optical_elements[k].type.type + ".")
331
332        # transform new rays back into "lab frame"
333        RayList = mgeo.TranslationRayList(RayList, -optical_elements[k].type.get_centre())
334        RayList = mgeo.RotationRayList(RayList, ex, mPrime)
335        RayList = mgeo.RotationRayList(RayList, ez, n)
336        RayList = mgeo.TranslationRayList(RayList, Position)
337
338        output_rays.append(RayList)
339
340    return output_rays

396def FindOptimalDistance(
397    Detector,
398    RayList: list[mray.Ray],
399    OptFor="intensity",
400    Amplitude: float = None,
401    Precision: int = 3,
402    IntensityWeighted=False,
403    verbose=False,
404):
405    """
406    Automatically finds the optimal the 'Detector'-distance for either
407    maximum intensity, or minimal spotsize or duration.
408
409    Parameters
410    ----------
411        Detector : Detector-object
412
413        RayList : list[Ray-objects]
414
415        OptFor : str, optional
416            "spotsize": minimizes the standard-deviation spot size *d* on the detector.
417            "duration": minimizes the standard-deviation *tau* of the ray-delays.
418            "intensity": Maximizes 1/tau/d^2.
419            Defaults to "intensity".
420
421        Amplitude : float, optional
422            The detector-distances within which the optimization is done will be
423            the distance of 'Detector' +/- Amplitude in mm.
424
425        Precision : int, optional
426            Precision-parameter for the search algorithm. For Precision = n,
427            it will iterate with search amplitudes decreasing until 10^{-n}.
428            Defaults to 3.
429
430        IntensityWeighted : bool, optional
431            Whether to weigh the calculation of spotsize and/or duration by the ray-intensities.
432            Defaults to False.
433
434        verbose : bool, optional
435            Whether to print results to the console.
436
437    Returns
438    -------
439        OptDetector : Detector-object
440            A new detector-instance with optimized distance.
441
442        OptSpotSize : float
443            Spotsize, calculated as standard deviation in mm of the spot-diagram
444            on the optimized detector.
445
446        OptDuration : float
447            Duration, calculated as standard deviation in fs of the ray-delays
448            on the optimized detector.
449    """
450
451    if OptFor not in ["intensity", "size", "duration"]:
452        raise NameError(
453            "I don`t recognize what you want to optimize the detector distance for. OptFor must be either 'intensity', 'size' or 'duration'."
454        )
455
456    FirstDistance = Detector.get_distance()
457    ListPointDetector2DCentre = Detector.get_PointList2DCentre(RayList)
458    SizeSpot = 2 * StandardDeviation(ListPointDetector2DCentre)
459    NumericalAperture = ReturnNumericalAperture(RayList, 1)
460    if Amplitude is None:
461        Amplitude = min(4 * np.ceil(SizeSpot / np.tan(np.arcsin(NumericalAperture))), FirstDistance)
462    # Step = Amplitude/10
463    Step = Amplitude / 5  # good enough in half the time ;-)
464
465    if verbose:
466        print(
467            f"Searching optimal detector position for *{OptFor}* within [{FirstDistance-Amplitude:.3f}, {FirstDistance+Amplitude:.3f}] mm...",
468            end="",
469            flush=True,
470        )
471    movingDetector = Detector.copy_detector()
472    for k in range(Precision + 1):
473        movingDetector, OptSpotSize, OptDuration = _FindOptimalDistanceBIS(
474            movingDetector, Amplitude * 0.1**k, Step * 0.1**k, RayList, OptFor, IntensityWeighted
475        )
476
477    if (
478        not FirstDistance - Amplitude + 10**-Precision
479        < movingDetector.get_distance()
480        < FirstDistance + Amplitude - 10**-Precision
481    ):
482        print("There`s no minimum-size/duration focus in the searched range.")
483
484    print(
485        "\r\033[K", end="", flush=True
486    )  # move to beginning of the line with \r and then delete the whole line with \033[K
487    return movingDetector, OptSpotSize, OptDuration

491def FindCentralRay(RayList: list[mray.Ray]):
492    """
493    Out of a the list of Ray-objects, RayList, return the one Ray-object whose Ray.number is 0.
494    This is the ray which the functions in ModuleSource set up as the central ray, i.e.
495    the one initially launched on the "optical axis".
496    Should there be several rays whose number-attribute =0 (which is a bad situation), only
497    the first one encountered in the list is returned.
498
499    Parameters
500    ----------
501        RayList : list of Ray-objects
502
503    Returns
504    -------
505        CentralRay : Ray-object
506    """
507    for k in RayList:
508        if k.number == 0:
509            return k
510    return None

513def StandardDeviation(List: list[float, np.ndarray]) -> float:
514    """
515    Return the standard deviation of elements in the supplied list.
516    If the elements are floats, it's simply the root-mean-square over the numbers.
517    If the elements are vectors (represented as numpy-arrays), the root-mean-square
518    of the distance of the points from their mean is calculated.
519
520    Parameters
521    ----------
522        List : list of floats or of np.ndarray
523            Numbers (in ART, typically delays)
524            or 2D or 3D vectors (in ART typically space coordinates)
525
526    Returns
527    -------
528        Std : float
529    """
530    if type(List[0]) in [int, float, np.float64]:
531        return np.std(List)
532    elif len(List[0]) > 1:
533        return np.sqrt(np.var(List, axis=0).sum())
534    else:
535        raise ValueError("StandardDeviation expects a list of floats or numpy-arrays as input, but got something else.")

538def WeightedStandardDeviation(List: list[float, np.ndarray], Weights: list[float]) -> float:
539    """
540    Return the weighted standard deviation of elements in the supplied list.
541    If the elements are floats, it's simply the root-mean-square over the weighted numbers.
542    If the elements are vectors (represented as numpy-arrays), the root-mean-square
543    of the weighted distances of the points from their mean is calculated.
544
545    Parameters
546    ----------
547        List : list of floats or of np.ndarray
548            Numbers (in ART, typically delays)
549            or 2D or 3D vectors (in ART typically space coordinates)
550
551        Weights : list of floats, same length as List
552            Parameters such as Ray.intensity to be used as weights
553
554    Returns
555    -------
556        Std : float
557    """
558    average = np.average(List, axis=0, weights=Weights, returned=False)
559    variance = np.average((List - average) ** 2, axis=0, weights=Weights, returned=False)
560    return np.sqrt(variance.sum())

564def ReturnNumericalAperture(RayList: list[mray.Ray], RefractiveIndex: float = 1) -> float:
565    r"""
566    Returns the numerical aperture associated with the supplied ray-bundle 'Raylist'.
567    This is $n\sin\theta$, where $\theta$ is the maximum angle between any of the rays and the central ray,
568    and $n$ is the refractive index of the propagation medium.
569
570    Parameters
571    ----------
572        RayList : list of Ray-object
573            The ray-bundle of which to determine the NA.
574
575        RefractiveIndex : float, optional
576            Refractive index of the propagation medium, defaults to =1.
577
578    Returns
579    -------
580        NA : float
581    """
582    CentralRay = FindCentralRay(RayList)
583    if CentralRay is None:
584        CentralVector = np.array([0, 0, 0])
585        for k in RayList:
586            CentralVector = CentralVector + k.vector
587        CentralVector = CentralVector / len(RayList)
588    else:
589        CentralVector = CentralRay.vector
590    ListAngleAperture = []
591    for k in RayList:
592        ListAngleAperture.append(mgeo.AngleBetweenTwoVectors(CentralVector, k.vector))
593
594    return np.sin(np.amax(ListAngleAperture)) * RefractiveIndex

598def ReturnAiryRadius(Wavelength: float, NumericalAperture: float) -> float:
599    r"""
600    Returns the radius of the Airy disk: $r = 1.22 \frac\{\lambda\}\{NA\}$,
601    i.e. the diffraction-limited radius of the focal spot corresponding to a given
602    numerical aperture $NA$ and a light wavelength $\lambda$.
603
604    For very small $NA<10^\{-3\}$, diffraction effects becomes negligible and the Airy Radius becomes meaningless,
605    so in that case, a radius of 0 is returned.
606
607    Parameters
608    ----------
609        Wavelength : float
610            Light wavelength in mm.
611
612        NumericalAperture : float
613
614    Returns
615    -------
616        AiryRadius : float
617    """
618    if NumericalAperture > 1e-3 and Wavelength is not None:
619        return 1.22 * 0.5 * Wavelength / NumericalAperture
620    else:
621        return 0  # for very small numerical apertures, diffraction effects becomes negligible and the Airy Radius becomes meaningless

640def save_compressed(obj, filename: str):
641    """Save (=pickle) an object 'obj' to a compressed file with name 'filename'."""
642    if not type(filename) == str:
643        filename = "kept_data_" + datetime.now().strftime("%Y-%m-%d-%Hh%M")
644
645    i = 0
646    while os.path.exists(filename + f"_{i}.xz"):
647        i += 1
648    filename = filename + f"_{i}"
649    # with gzip.open(filename + '.gz', 'wb') as f:
650    with lzma.open(filename + ".xz", "wb") as f:
651        pickle.dump(obj, f)
652    print("Saved results to " + filename + ".xz.")
653    print("->To reload from disk do: kept_data = mp.load_compressed('" + filename + "')")

656def load_compressed(filename: str):
657    """Load (=unpickle) an object 'obj' from a compressed file with name 'filename'."""
658    # with gzip.open(filename + '.gz', 'rb') as f:
659    with lzma.open(filename + ".xz", "rb") as f:
660        obj = pickle.load(f)
661    return obj