Merge pull request #303 from AstarVienna/fh/zorder

Some refactoring of z_order fuctionality

scopesim/effects/effects_utils.py

@@ -1,7 +1,10 @@
 """TBA."""
 
+import logging
 import inspect
 from copy import deepcopy, copy
+from collections.abc import Iterable
+
 from astropy.table import Table
 
 from .. import effects as efs
@@ -107,3 +110,29 @@ def scopesim_effect_classes(base_effect=efs.Effect):
     sorted_effects = {key: efs_dict[key] for key in sorted(efs_dict)}
 
     return sorted_effects
+
+
+def z_order_in_range(z_eff, z_range: range) -> bool:
+    """
+    Return True if any of the z_orders in `z_eff` is in the given range.
+
+    The `z_range` parameter can be constructed as ``range(z_min, z_max)``.
+
+    Parameters
+    ----------
+    z_eff : int or list of ints
+        z_order(s) of the effect.
+    z_range : range
+        range object of allowed z_order values.
+
+    Returns
+    -------
+    bool
+        True if at least one z_order is in range, False otherwise.
+
+    """
+    if not isinstance(z_eff, Iterable):
+        logging.warning("z_order %d should be a single-item iterable", z_eff)
+        z_eff = [z_eff]
+
+    return any(zi in z_range for zi in z_eff)

scopesim/optics/optical_element.py

@@ -6,7 +6,8 @@
 from astropy.table import Table
 
 from .. import effects as efs
-from ..effects.effects_utils import make_effect, get_all_effects
+from ..effects.effects_utils import (make_effect, get_all_effects,
+                                     z_order_in_range)
 from ..utils import write_report
 from ..reports.rst_utils import table_to_rst
 from .. import rc
@@ -88,40 +89,70 @@ def add_effect(self, effect):
     def get_all(self, effect_class):
         return get_all_effects(self.effects, effect_class)
 
-    def get_z_order_effects(self, z_level):
-        if isinstance(z_level, int):
-            zmin = z_level
-            zmax = zmin + 99
-        elif isinstance(z_level, (tuple, list)):
-            zmin, zmax = z_level[:2]
+    def get_z_order_effects(self, z_level: int, z_max: int = None):
+        """
+        Yield all effects in the given 100-range of `z_level`.
+
+        E.g., ``z_level=200`` will yield all effect with a z_order between
+        200 and 299. Optionally, the upper limit can be set manually with the
+        optional argument `z_max`.
+
+        Parameters
+        ----------
+        z_level : int
+            100-range of z_orders.
+        z_max : int, optional
+            Optional upper bound. This is currently not used anywhere in
+            ScopeSim, but the functionality is tested. If None (default), this
+            will be set to ``z_level + 99``.
+
+        Raises
+        ------
+        TypeError
+            Raised if either `z_level` or `z_max` is not of int type.
+        ValueError
+            Raised if `z_max` (if given) is less than `z_level`.
+
+        Yields
+        ------
+        eff : Iterator of effects
+            Iterator containing all effect objects in the given z_order range.
+
+        """
+        if not isinstance(z_level, int):
+            raise TypeError(f"z_level must be int, got {type(z_level)=}")
+        if z_max is not None and not isinstance(z_max, int):
+            raise TypeError(f"If given, z_max must be int, got {type(z_max)=}")
+
+        z_min = z_level
+        if z_max is not None:
+            if z_max < z_min:
+                raise ValueError(
+                    "z_max must be greater (or equal to) z_level, but "
+                    f"{z_max=} < {z_level=}.")
         else:
-            zmin, zmax = 0, 999
+            z_max = z_min + 100  # range doesn't include final element -> 100
+        z_range = range(z_min, z_max)
 
-        effects = []
         for eff in self.effects:
-            if eff.include and "z_order" in eff.meta:
-                z = eff.meta["z_order"]
-                if isinstance(z, (list, tuple)):
-                    if any(zmin <= zi <= zmax for zi in z):
-                        effects.append(eff)
-                else:
-                    if zmin <= z <= zmax:
-                        effects.append(eff)
+            if not eff.include or "z_order" not in eff.meta:
+                continue
 
-        return effects
+            if z_order_in_range(eff.meta["z_order"], z_range):
+                yield eff
+
+    def _get_matching_effects(self, effect_classes):
+        return (eff for eff in self.effects if isinstance(eff, effect_classes))
 
     @property
     def surfaces_list(self):
-        _ter_list = [effect for effect in self.effects
-                     if isinstance(effect, (efs.SurfaceList, efs.FilterWheel,
-                                            efs.TERCurve))]
-        return _ter_list
+        effect_classes = (efs.SurfaceList, efs.FilterWheel, efs.TERCurve)
+        return list(self._get_matching_effects(effect_classes))
 
     @property
     def masks_list(self):
-        _mask_list = [effect for effect in self.effects if
-                      isinstance(effect, (efs.ApertureList, efs.ApertureMask))]
-        return _mask_list
+        effect_classes = (efs.ApertureList, efs.ApertureMask)
+        return list(self._get_matching_effects(effect_classes))
 
     def list_effects(self):
         elements = [self.meta["name"]] * len(self.effects)
@@ -223,15 +254,18 @@ def _repr_pretty_(self, p, cycle):
 
     @property
     def properties_str(self):
-        prop_str = ""
-        max_key_len = max(len(key) for key in self.properties.keys())
-        padlen = max_key_len + 4
-        for key in self.properties:
-            if key not in {"comments", "changes", "description", "history",
-                           "report"}:
-                prop_str += f"{key:>{padlen}} : {self.properties[key]}\n"
-
-        return prop_str
+        # TODO: This seems to be used only in the report below.
+        #       Once the report uses stream writing, change this to a function
+        #       that simply write to that same stream...
+        padlen = max(len(key) for key in self.properties) + 4
+        exclude = {"comments", "changes", "description", "history", "report"}
+
+        with StringIO() as str_stream:
+            for key in self.properties.keys() - exclude:
+                str_stream.write(f"{key:>{padlen}} : {self.properties[key]}\n")
+            output = str_stream.getvalue()
+
+        return output
 
     def report(self, filename=None, output="rst", rst_title_chars="^#*+",
                **kwargs):

scopesim/optics/optics_manager.py

@@ -139,7 +139,7 @@ def get_all(self, class_type):
 
         return effects
 
-    def get_z_order_effects(self, z_level):
+    def get_z_order_effects(self, z_level: int):
         """
         Return a list of all effects with a z_order keywords within `z_level`.
 
@@ -150,82 +150,98 @@ def get_z_order_effects(self, z_level):
         - Apply Source altering effects - z_order = 200..299
         - Apply FOV specific (3D) effects - z_order = 300..399
         - Apply FOV-independent (2D) effects - z_order = 400..499
+        - Apply XXX effects - z_order = 500..599
+        - Apply XXX effects - z_order = 600..699
+        - Apply lambda-independent 2D image plane effects - z_order = 700..799
+        - Apply detector effects - z_order = 800..899
+        - Apply detector array effects - z_order = 900..999
 
         Parameters
         ----------
-        z_level : int, tuple
-            [0, 100, 200, 300, 400, 500]
+        z_level : {0, 100, 200, 300, 400, 500, 600, 700, 800, 900}
+            100-range of z_orders.
 
         Returns
         -------
         effects : list of Effect objects
 
         """
-        effects = []
-        for opt_el in self.optical_elements:
-            effects += opt_el.get_z_order_effects(z_level)
+        def _gather_effects():
+            for opt_el in self.optical_elements:
+                yield from opt_el.get_z_order_effects(z_level)
 
-        return effects
+        def _sortkey(eff):
+            return next(z % 100 for z in eff.meta["z_order"] if z >= z_level)
+
+        # return sorted(_gather_effects(), key=_sortkey)
+        return list(_gather_effects())
 
     @property
     def is_spectroscope(self):
+        """Return True if any of the effects is a spectroscope."""
         return is_spectroscope(self.all_effects)
 
     @property
     def image_plane_headers(self):
+        """Get headers from detector setup effects."""
         detector_lists = self.detector_setup_effects
-        headers = [det_list.image_plane_header for det_list in detector_lists]
-
         if not detector_lists:
-            raise ValueError(f"No DetectorList objects found. {detector_lists}")
+            raise ValueError("No DetectorList objects found.")
 
-        return headers
+        return [det_list.image_plane_header for det_list in detector_lists]
 
     @property
     def detector_array_effects(self):
+        """Get effects with z_order = 900...999."""
         return self.get_z_order_effects(900)
 
     @property
     def detector_effects(self):
+        """Get effects with z_order = 800...899."""
         return self.get_z_order_effects(800)
 
     @property
     def image_plane_effects(self):
-        effects = self.get_z_order_effects(700)
-        return effects
+        """Get effects with z_order = 700...799."""
+        return self.get_z_order_effects(700)
 
     @property
     def fov_effects(self):
-        effects = self.get_z_order_effects(600)
-        return effects
+        """Get effects with z_order = 600...699."""
+        return self.get_z_order_effects(600)
 
     @property
     def source_effects(self):
+        """Get effects with z_order = 500...599."""
         return self.get_z_order_effects(500)   # Transmission
 
     @property
     def detector_setup_effects(self):
-        # !!! Only DetectorLists go in here !!!
+        """Get effects with z_order = 400...499 (DetectorLists only!)."""
         return self.get_z_order_effects(400)
 
     @property
     def image_plane_setup_effects(self):
+        """Get effects with z_order = 300...399."""
         return self.get_z_order_effects(300)
 
     @property
     def fov_setup_effects(self):
+        """Get effects with z_order = 200...299."""
         # Working out where to set wave_min, wave_max
         return self.get_z_order_effects(200)
 
     @property
     def surfaces_table(self):
+        """Get combined surface table from effects with z_order = 100...199."""
         if self._surfaces_table is None:
             surface_like_effects = self.get_z_order_effects(100)
             self._surfaces_table = combine_surface_effects(surface_like_effects)
         return self._surfaces_table
 
     @property
     def all_effects(self):
+        """Get all effects in all optical elements."""
         return [eff for opt_eff in self.optical_elements for eff in opt_eff]
 
     @property

scopesim/tests/tests_optics/test_OpticalElement.py

@@ -52,11 +52,13 @@ def test_currsys_ignore_effects_have_false_include_flag(self, atmo_yaml_dict):
 
 @pytest.mark.usefixtures("patch_mock_path")
 class TestOpticalElementGetZOrderEffects:
-    @pytest.mark.parametrize("z_orders, n", [(0, 2), (100, 1), ([200, 299], 1)])
-    def test_returns_the_effects_with_z_values(self, z_orders, n,
+    @pytest.mark.parametrize("z_lvl, zmax, n", [(0, None, 2),
+                                                (100, None, 1),
+                                                (200, 299, 1)])
+    def test_returns_the_effects_with_z_values(self, z_lvl, zmax, n,
                                                detector_yaml_dict):
         opt_el = opt_elem.OpticalElement(detector_yaml_dict)
-        assert len(opt_el.get_z_order_effects(z_orders)) == n
+        assert len(list(opt_el.get_z_order_effects(z_lvl, zmax))) == n
 
 
 @pytest.mark.usefixtures("patch_mock_path")