ModuleOpticalChain

 27class OpticalChain:
 28    """
 29    The OpticalChain represents the whole optical setup to be simulated:
 30    Its main attributes are a list source-[Rays](ModuleOpticalRay.html) and
 31    a list of successive [OpticalElements](ModuleOpticalElement.html).
 32
 33    The method OpticalChain.get_output_rays() returns an associated list of lists of
 34    [Rays](ModuleOpticalRay.html), each calculated by ray-tracing from one
 35    [OpticalElement](ModuleOpticalElement.html) to the next.
 36    So OpticalChain.get_output_rays()[i] is the bundle of [Rays](ModuleOpticalRay.html) *after*
 37    optical_elements[i].
 38
 39    The string "description" can contain a short description of the optical setup, or similar notes.
 40
 41    The OpticalChain can be visualized quickly with the method OpticalChain.quickshow(),
 42    and more nicely with OpticalChain.render().
 43
 44    The class also provides methods for (mis-)alignment of the source-[Ray](ModuleOpticalRay.html)-bundle and the
 45    [OpticalElements](ModuleOpticalElement.html), as well as methods for producing
 46    a list of OpticalChain-objects containing variations of itself.
 47
 48    Attributes
 49    ----------
 50        source_rays : list[mray.Ray]
 51            List of source rays, which are to be traced.
 52
 53        optical_elements : list[moe.OpticalElement]
 54            List of successive optical elements.
 55
 56        description : str
 57            A string to describe the optical setup.
 58
 59        loop_variable_name : str
 60            A string naming a parameter that is varied in a list of OpticalChain-objects,
 61            which is useful when looping over variations of an initial configuration.
 62
 63        loop_variable_value : float
 64            The value of that varied parameter, which is useful when looping over
 65            variations of an initial configuration.
 66
 67    Methods
 68    ----------
 69        copy_chain()
 70
 71        get_output_rays()
 72
 73        quickshow()
 74
 75        render()
 76
 77        ----------
 78
 79        shift_source(axis, distance)
 80
 81        tilt_source(self, axis, angle)
 82
 83        get_source_loop_list(axis, loop_variable_values)
 84
 85        ----------
 86
 87        rotate_OE(OEindx, axis, angle)
 88
 89        shift_OE(OEindx, axis, distance)
 90
 91        get_OE_loop_list(OEindx, axis, loop_variable_values)
 92
 93    """
 94
 95    def __init__(
 96        self, source_rays, optical_elements, description="", loop_variable_name=None, loop_variable_value=None
 97    ):
 98        """
 99        Parameters
100        ----------
101            source_rays : list[mray.Ray]
102                List of source rays, which are to be traced.
103
104            optical_elements : list[moe.OpticalElement]
105                List of successive optical elements.
106
107            description : str, optional
108                A string to describe the optical setup. Defaults to ''.
109
110            loop_variable_name : str, optional
111                A string naming a parameter that is varied in a list of OpticalChain-objects.
112                Defaults to None.
113
114            loop_variable_value : float
115                The value of that varied parameter, which is useful when looping over
116                variations of an initial configuration. Defaults to None.
117        """
118        self.source_rays = copy.deepcopy(source_rays)
119        # deepcopy so this object doesn't get changed when the global source_rays changes "outside"
120        self.optical_elements = copy.deepcopy(optical_elements)
121        # deepcopy so this object doesn't get changed when the global optical_elements changes "outside"
122        self.description = description
123        self.loop_variable_name = loop_variable_name
124        self.loop_variable_value = loop_variable_value
125        self._output_rays = None
126        self._last_optical_elements_hash = None  # for now we don't care which element was changed.
127        # we just always do the whole raytracing again
128        self._last_source_rays_hash = None
129
130    # using property decorator
131    # a getter function
132    @property
133    def source_rays(self):
134        return self._source_rays
135
136    # a setter function
137    @source_rays.setter
138    def source_rays(self, source_rays):
139        if type(source_rays) == list and all(isinstance(x, mray.Ray) for x in source_rays):
140            self._source_rays = source_rays
141        else:
142            raise TypeError("Source_rays must be list of Ray-objects.")
143
144    @property
145    def optical_elements(self):
146        return self._optical_elements
147
148    @optical_elements.setter
149    def optical_elements(self, optical_elements):
150        if type(optical_elements) == list and all(isinstance(x, moe.OpticalElement) for x in optical_elements):
151            self._optical_elements = optical_elements
152        else:
153            raise TypeError("Optical_elements must be list of OpticalElement-objects.")
154
155    @property
156    def loop_variable_name(self):
157        return self._loop_variable_name
158
159    @loop_variable_name.setter
160    def loop_variable_name(self, loop_variable_name):
161        if type(loop_variable_name) == str or (loop_variable_name is None):
162            self._loop_variable_name = loop_variable_name
163        else:
164            raise TypeError("loop_variable_name must be a string.")
165
166    @property
167    def loop_variable_value(self):
168        return self._loop_variable_value
169
170    @loop_variable_value.setter
171    def loop_variable_value(self, loop_variable_value):
172        if type(loop_variable_value) in [int, float, np.float64] or (loop_variable_value is None):
173            self._loop_variable_value = loop_variable_value
174        else:
175            raise TypeError("loop_variable_value must be a number of types int or float.")
176
177    # %% METHODS ##################################################################
178
179    def copy_chain(self):
180        """Return another optical chain with the same source, optical elements and description-string as this one."""
181        return OpticalChain(self.source_rays, self.optical_elements, self.description)
182
183    def get_output_rays(self, **kwargs):
184        """
185        Returns the list of (lists of) output rays, calculate them if this hasn't been done yet,
186        or if the source-ray-bundle or anything about the optical elements has changed.
187        """
188        current_source_rays_hash = mp._hash_list_of_objects(self.source_rays)
189        current_optical_elements_hash = mp._hash_list_of_objects(self.optical_elements)
190        if (current_source_rays_hash != self._last_source_rays_hash) or (
191            current_optical_elements_hash != self._last_optical_elements_hash
192        ):
193            print("...ray-tracing...", end="", flush=True)
194            self._output_rays = mp.RayTracingCalculation(self.source_rays, self.optical_elements, **kwargs)
195            print(
196                "\r\033[K", end="", flush=True
197            )  # move to beginning of the line with \r and then delete the whole line with \033[K
198
199            self._last_source_rays_hash = current_source_rays_hash
200            self._last_optical_elements_hash = current_optical_elements_hash
201
202        return self._output_rays
203
204    def quickshow(self):
205        """Render an image of the optical chain it with settings that prioritize
206        speed over great looks. This lets the user quickly visualize their
207        optical setup to check if all the angles are set as they want."""
208        maxRays = 30
209        maxOEpoints = 1500
210        QuickOpticalChain = self.copy_chain()
211        QuickOpticalChain.source_rays = np.random.choice(self.source_rays, maxRays, replace=False).tolist()
212        quickfig = mplots.RayRenderGraph(QuickOpticalChain, None, maxRays, maxOEpoints)
213        return quickfig
214
215    def render(self):
216        """Create a fairly good-looking 3D rendering of the optical chain."""
217        maxRays = 150
218        maxOEpoints = 3000
219        fig = mplots.RayRenderGraph(self, None, maxRays, maxOEpoints)
220        return fig
221
222    # %%  methods to (mis-)align the optical chain; just uses the corresponding methods of the OpticalElement class...
223
224    def shift_source(self, axis: (str, np.ndarray), distance: float):
225        """
226        Shift source ray bundle by distance (in mm) along the 'axis' specified as
227        a lab-frame vector (numpy-array of length 3) or as one of the strings
228        "vert", "horiz", or "random".
229
230        In the latter case, the reference is the incidence plane of the first
231        non-normal-incidence mirror after the source. If there is none, you will
232        be asked to rather specify the axis as a 3D-numpy-array.
233
234        axis = "vert" means the source position is shifted along the axis perpendicular
235        to that incidence plane, i.e. "vertically" away from the former incidence plane..
236
237        axis = "horiz" means the source direciton is rotated about an axis in that
238        incidence plane and perpendicular to the current source direction,
239        i.e. "horizontally" in the incidence plane, but retaining the same distance
240        of source and first optical element.
241
242        axis = "random" means the the source direction shifted in a random direction
243        within in the plane perpendicular to the current source direction,
244        e.g. simulating a fluctuation of hte transverse source position.
245
246        Parameters
247        ----------
248            axis : np.ndarray or str
249                Shift axis, specified either as a 3D lab-frame vector or as one
250                of the strings "vert", "horiz", or "random".
251
252            distance : float
253                Shift distance in mm.
254
255        Returns
256        -------
257            Nothing, just modifies the property 'source_rays'.
258        """
259        if type(distance) not in [int, float, np.float64]:
260            raise ValueError('The "distance"-argument must be an int or float number.')
261
262        central_ray_vector = mp.FindCentralRay(self.source_rays).vector
263        mirror_indcs = [i for i, OE in enumerate(self.optical_elements) if "Mirror" in OE.type.type]
264
265        OEnormal = None
266        for i in mirror_indcs:
267            ith_OEnormal = self.optical_elements[i].normal
268            if np.linalg.norm(np.cross(central_ray_vector, ith_OEnormal)) > 1e-10:
269                OEnormal = ith_OEnormal
270                break
271        if OEnormal is None:
272            raise Exception(
273                "There doesn't seem to be a non-normal-incidence mirror in this optical chain, \
274                            so you should rather give 'axis' as a numpy-array of length 3."
275            )
276
277        if type(axis) == np.ndarray and len(axis) == 3:
278            translation_vector = axis
279        else:
280            perp_axis = np.cross(central_ray_vector, OEnormal)
281            horiz_axis = np.cross(perp_axis, central_ray_vector)
282
283            if axis == "vert":
284                translation_vector = perp_axis
285            elif axis == "horiz":
286                translation_vector = horiz_axis
287            elif axis == "random":
288                translation_vector = (
289                    np.random.uniform(low=-1, high=1, size=1) * perp_axis
290                    + np.random.uniform(low=-1, high=1, size=1) * horiz_axis
291                )
292            else:
293                raise ValueError(
294                    'The shift direction must be specified by "axis" as one of ["vert", "horiz", "random"].'
295                )
296
297        self.source_rays = mgeo.TranslationRayList(self.source_rays, distance * mgeo.Normalize(translation_vector))
298
299    def tilt_source(self, axis: (str, np.ndarray), angle: float):
300        """
301        Rotate source ray bundle by angle around an axis, specified as
302        a lab-frame vector (numpy-array of length 3) or as one of the strings
303        "in_plane", "out_plane" or "random" direction.
304
305        In the latter case, the function considers the incidence plane of the first
306        non-normal-incidence mirror after the source. If there is none, you will
307        be asked to rather specify the axis as a 3D-numpy-array.
308
309        axis = "in_plane" means the source direction is rotated about an axis
310        perpendicular to that incidence plane, which tilts the source
311        "horizontally" in the same plane.
312
313        axis = "out_plane" means the source direciton is rotated about an axis
314        in that incidence plane and perpendicular to the current source direction,
315        which tilts the source "vertically" out of the former incidence plane.
316
317        axis = "random" means the the source direction is tilted in a random direction,
318        e.g. simulating a beam pointing fluctuation.
319
320        Attention, "angle" is given in deg, so as to remain consitent with the
321        conventions of other functions, although pointing is mostly talked about
322        in mrad instead.
323
324        Parameters
325        ----------
326            axis : np.ndarray or str
327                Shift axis, specified either as a 3D lab-frame vector or as one
328                of the strings "in_plane", "out_plane", or "random".
329
330            angle : float
331                Rotation angle in degree.
332
333        Returns
334        -------
335            Nothing, just modifies the property 'source_rays'.
336        """
337        if type(angle) not in [int, float, np.float64]:
338            raise ValueError('The "angle"-argument must be an int or float number.')
339
340        central_ray_vector = mp.FindCentralRay(self.source_rays).vector
341        mirror_indcs = [i for i, OE in enumerate(self.optical_elements) if "Mirror" in OE.type.type]
342
343        OEnormal = None
344        for i in mirror_indcs:
345            ith_OEnormal = self.optical_elements[i].normal
346            if np.linalg.norm(np.cross(central_ray_vector, ith_OEnormal)) > 1e-10:
347                OEnormal = ith_OEnormal
348                break
349        if OEnormal is None:
350            raise Exception(
351                "There doesn't seem to be a non-normal-incidence mirror in this optical chain, \
352                            so you should rather give 'axis' as a numpy-array of length 3."
353            )
354
355        if type(axis) == np.ndarray and len(axis) == 3:
356            rot_axis = axis
357        else:
358            rot_axis_in = np.cross(central_ray_vector, OEnormal)
359            rot_axis_out = np.cross(rot_axis_in, central_ray_vector)
360            if axis == "in_plane":
361                rot_axis = rot_axis_in
362            elif axis == "out_plane":
363                rot_axis = rot_axis_out
364            elif axis == "random":
365                rot_axis = (
366                    np.random.uniform(low=-1, high=1, size=1) * rot_axis_in
367                    + np.random.uniform(low=-1, high=1, size=1) * rot_axis_out
368                )
369            else:
370                raise ValueError(
371                    'The tilt axis must be specified by as one of ["in_plane", "out_plane", "random"] or as a numpy-array of length 3.'
372                )
373
374        self.source_rays = mgeo.RotationAroundAxisRayList(self.source_rays, rot_axis, np.deg2rad(angle))
375
376    def get_source_loop_list(self, axis: str, loop_variable_values: np.ndarray):
377        """
378        Produces a list of OpticalChain-objects, which are all variations of this
379        instance by moving the source-ray-bundle.
380        The variation is specified by axis as one of
381        ["tilt_in_plane", "tilt_out_plane", "tilt_random", "shift_vert", "shift_horiz", "shift_random"],
382        by the values given in the list or numpy-array "loop_variable_values", e.g. np.linspace(start, stop, number).
383        This list can then be looped over by ARTmain.
384
385        Parameters
386        ----------
387            axis : np.ndarray or str
388                Shift/Rotation axis for the source-modification, specified either
389                as a 3D lab-frame vector or as one of the strings
390                ["tilt_in_plane", "tilt_out_plane", "tilt_random",\
391                 "shift_vert", "shift_horiz", "shift_random"].
392
393             loop_variable_values : list or np.ndarray
394                Values of the shifts (mm) or rotations (deg).
395
396        Returns
397        -------
398            OpticalChainList : list[OpticalChain]
399        """
400        if axis not in [
401            "tilt_in_plane",
402            "tilt_out_plane",
403            "tilt_random",
404            "shift_vert",
405            "shift_horiz",
406            "shift_random",
407            "all_random",
408        ]:
409            raise ValueError(
410                'For automatic loop-list generation, the axis must be one of ["tilt_in_plane", "tilt_out_plane", "tilt_random", "shift_vert", "shift_horiz", "shift_random"].'
411            )
412        if type(loop_variable_values) not in [list, np.ndarray]:
413            raise ValueError(
414                "For automatic loop-list generation, the loop_variable_values must be a list or a numpy-array."
415            )
416
417        loop_variable_name_strings = {
418            "tilt_in_plane": "source tilt in-plane (deg)",
419            "tilt_out_plane": "source tilt out-of-plane (deg)",
420            "tilt_random": "source tilt random axis (deg)",
421            "shift_vert": "source shift vertical (mm)",
422            "shift_horiz": "source shift horizontal (mm)",
423            "shift_random": "source shift random-direction (mm)",
424        }
425        loop_variable_name = loop_variable_name_strings[axis]
426
427        OpticalChainList = []
428        for x in loop_variable_values:
429            # always start with a fresh deep-copy the AlignedOpticalChain, to then modify it and append it to the list
430            ModifiedOpticalChain = self.copy_chain()
431            ModifiedOpticalChain.loop_variable_name = loop_variable_name
432            ModifiedOpticalChain.loop_variable_value = x
433
434            if axis in ["tilt_in_plane", "tilt_out_plane", "tilt_random"]:
435                ModifiedOpticalChain.tilt_source(axis[5:], x)
436            elif axis in ["shift_vert", "shift_horiz", "shift_random"]:
437                ModifiedOpticalChain.shift_source(axis[6:], x)
438
439            # append the modified optical chain to the list
440            OpticalChainList.append(ModifiedOpticalChain)
441
442        return OpticalChainList
443
444    # %%
445    def rotate_OE(self, OEindx: int, axis: str, angle: float):
446        """
447        Rotate the optical element OpticalChain.optical_elements[OEindx] about
448        axis specified by "pitch", "roll", "yaw", or "random" by angle in degrees.
449
450        Parameters
451        ----------
452            OEindx : int
453                Index of the optical element to modify out of OpticalChain.optical_elements.
454
455            axis : str
456                Rotation axis, specified as one of the strings
457                "pitch", "roll", "yaw", or "random".
458                These define the rotations as specified for the corresponding methods
459                of the OpticalElement-class.
460
461            angle : float
462                Rotation angle in degree.
463
464        Returns
465        -------
466            Nothing, just modifies OpticalChain.optical_elements[OEindx].
467        """
468        if abs(OEindx) > len(self.optical_elements):
469            raise ValueError(
470                'The "OEnumber"-argument is out of range compared to the length of OpticalChain.optical_elements.'
471            )
472        if type(angle) not in [int, float, np.float64]:
473            raise ValueError('The "angle"-argument must be an int or float number.')
474
475        if axis == "pitch":
476            self.optical_elements[OEindx].rotate_pitch_by(angle)
477        elif axis == "roll":
478            self.optical_elements[OEindx].rotate_roll_by(angle)
479        elif axis == "yaw":
480            self.optical_elements[OEindx].rotate_yaw_by(angle)
481        elif axis in ("random", "rotate_random"):
482            self.optical_elements[OEindx].rotate_random_by(angle)
483        else:
484            raise ValueError('The "axis"-argument must be a string out of ["pitch", "roll", "yaw", "random"].')
485
486    def shift_OE(self, OEindx: int, axis: str, distance: float):
487        """
488        Shift the optical element OpticalChain.optical_elements[OEindx] along
489        axis specified by "normal", "major", "cross", or "random" by distance in mm.
490
491        Parameters
492        ----------
493            OEindx : int
494                Index of the optical element to modify out of OpticalChain.optical_elements.
495
496            axis : str
497                Rotation axis, specified as one of the strings
498                "normal", "major", "cross", or "random".
499                These define the shifts as specified for the corresponding methods
500                of the OpticalElement-class.
501
502            distance : float
503                Rotation angle in degree.
504
505        Returns
506        -------
507            Nothing, just modifies OpticalChain.optical_elements[OEindx].
508        """
509        if abs(OEindx) > len(self.optical_elements):
510            raise ValueError(
511                'The "OEnumber"-argument is out of range compared to the length of OpticalChain.optical_elements.'
512            )
513        if type(distance) not in [int, float, np.float64]:
514            raise ValueError('The "dist"-argument must be an int or float number.')
515
516        if axis == "normal":
517            self.optical_elements[OEindx].shift_along_normal(distance)
518        elif axis == "major":
519            self.optical_elements[OEindx].shift_along_major(distance)
520        elif axis == "cross":
521            self.optical_elements[OEindx].shift_along_cross(distance)
522        elif axis == "random":
523            self.optical_elements[OEindx].shift_along_random(distance)
524        else:
525            raise ValueError('The "axis"-argument must be a string out of ["normal", "major", "cross", "random"].')
526
527        # some function that randomly misalings one, or several or all ?
528
529    def get_OE_loop_list(self, OEindx: int, axis: str, loop_variable_values: np.ndarray):
530        """
531        Produces a list of OpticalChain-objects, which are all variations of
532        this instance by moving one degree of freedom of its optical element
533        with index OEindx.
534        The vaiations is specified by 'axis' as one of
535        ["pitch", "roll", "yaw", "rotate_random",\
536         "shift_normal", "shift_major", "shift_cross", "shift_random"],
537        by the values given in the list or numpy-array "loop_variable_values",
538        e.g. np.linspace(start, stop, number).
539        This list can then be looped over by ARTmain.
540
541        Parameters
542        ----------
543            OEindx :int
544                Index of the optical element to modify out of OpticalChain.optical_elements.
545
546            axis : np.ndarray or str
547                Shift/Rotation axis, specified as one of the strings
548                ["pitch", "roll", "yaw", "rotate_random",\
549                 "shift_normal", "shift_major", "shift_cross", "shift_random"].
550
551             loop_variable_values : list or np.ndarray
552                Values of the shifts (mm) or rotations (deg).
553
554        Returns
555        -------
556            OpticalChainList : list[OpticalChain]
557
558        """
559        if abs(OEindx) > len(self.optical_elements):
560            raise ValueError(
561                'The "OEnumber"-argument is out of range compared to the length of OpticalChain.optical_elements.'
562            )
563        if axis not in [
564            "pitch",
565            "roll",
566            "yaw",
567            "rotate_random",
568            "shift_normal",
569            "shift_major",
570            "shift_cross",
571            "shift_random",
572        ]:
573            raise ValueError(
574                'For automatic loop-list generation, the axis must be one of ["pitch", "roll", "yaw", "shift_normal", "shift_major", "shift_cross"].'
575            )
576        if type(loop_variable_values) not in [list, np.ndarray]:
577            raise ValueError(
578                "For automatic loop-list generation, the loop_variable_values must be a list or a numpy-array."
579            )
580
581        OE_name = self.optical_elements[OEindx].type.type + "_idx_" + str(OEindx)
582
583        loop_variable_name_strings = {
584            "pitch": OE_name + " pitch rotation (deg)",
585            "roll": OE_name + " roll rotation (deg)",
586            "yaw": OE_name + " yaw rotation (deg)",
587            "rotate_random": OE_name + " random rotation (deg)",
588            "shift_normal": OE_name + " shift along normal axis (mm)",
589            "shift_major": OE_name + " shift along major axis (mm)",
590            "shift_cross": OE_name + " shift along (normal x major)-direction (mm)",
591            "shift_random": OE_name + " shift along random axis (mm)",
592        }
593        loop_variable_name = loop_variable_name_strings[axis]
594
595        OpticalChainList = []
596        for x in loop_variable_values:
597            # always start with a fresh deep-copy the AlignedOpticalChain, to then modify it and append it to the list
598            ModifiedOpticalChain = self.copy_chain()
599            ModifiedOpticalChain.loop_variable_name = loop_variable_name
600            ModifiedOpticalChain.loop_variable_value = x
601
602            if axis in ("pitch", "roll", "yaw", "rotate_random"):
603                ModifiedOpticalChain.rotate_OE(OEindx, axis, x)
604            elif axis in ("shift_normal", "shift_major", "shift_cross", "shift_random"):
605                ModifiedOpticalChain.shift_OE(OEindx, axis[6:], x)
606
607            # append the modified optical chain to the list
608            OpticalChainList.append(ModifiedOpticalChain)
609
610        return OpticalChainList