PupilMaskWheel

scopesim/effects/ter_curves.py

@@ -896,9 +896,80 @@ def __init__(self, transmission, cmds=None, **kwargs):
                          emissivity=[0., 0.], **self.params)
 
     def update_transmission(self, transmission, **kwargs):
+        """Set a new transmission value"""
         self.__init__(transmission, **kwargs)
 
 
+class PupilMaskWheel(Effect):
+    """
+    Wheel holding a selection of predefined pupil masks
+
+    Currently, changing the PupilMask only changes the relative transmission.
+    Eventually, the change should also affect the total PSF of the OpticalTrain.
+
+    Example
+    -------
+    ::
+
+       name : pupil_masks
+       class : PupilMaskWheel
+       kwargs:
+           pupil_masks: [("name1", transmission1),
+                         ("name2", transmission2)]
+           current_mask: "open"
+    """
+    required_keys = {"pupil_masks", "current_mask"}
+    z_order: ClassVar[tuple[int, ...]] = (126, 226, 526)
+    report_plot_include: ClassVar[bool] = False
+    report_table_include: ClassVar[bool] = True
+    report_table_rounding: ClassVar[int] = 4
+    _current_str = "current_mask"
+
+    def __init__(self, cmds=None, **kwargs):
+        super().__init__(cmds=cmds, **kwargs)
+        check_keys(kwargs, self.required_keys, action="error")
+
+        self.meta.update(kwargs)
+        mask_dict = from_currsys(self.meta["pupil_masks"], cmds=self.cmds)
+        names = mask_dict['names']
+        transmissions = mask_dict['transmissions']
+        self.masks = {}
+        for name, trans in zip(names, transmissions):
+            kwargs["name"] = name
+            self.masks[name] = PupilTransmission(transmission=trans,
+                                                 cmds=cmds,
+                                                 **kwargs)
+        self.table = self.get_table(mask_dict)
+
+    def apply_to(self, obj, **kwargs):
+        """Use ``apply_to`` of current pupil mask."""
+        return self.current_mask.apply_to(obj, **kwargs)
+
+    def change_mask(self, maskname=None):
+        """Change the current pupil mask."""
+        if not maskname or maskname in self.masks:
+            self.meta["current_mask"] = maskname
+            self.include = maskname
+        else:
+            raise ValueError(f"Unknown pupil mask requested: {maskname}")
+
+    @property
+    def current_mask(self):
+        """Return the currently used pupil mask."""
+        currmask = from_currsys(self.meta['current_mask'], cmds=self.cmds)
+        if not currmask:
+            return False
+        return self.masks[currmask]
+
+    def __getattr__(self, item):
+        return getattr(self.current_mask, item)
+
+    def get_table(self, maskdict):
+        """Create a table of pupil masks with throughput."""
+        tbl = Table(maskdict)
+        return tbl
+
+
 class ADCWheel(Effect):
     """
     Wheel holding a selection of predefined atmospheric dispersion correctors.