diff --git a/scopesim/effects/effects_utils.py b/scopesim/effects/effects_utils.py index 58751b9f..780c1bbc 100644 --- a/scopesim/effects/effects_utils.py +++ b/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) diff --git a/scopesim/optics/optical_element.py b/scopesim/optics/optical_element.py index 11f4dc5e..f234a901 100644 --- a/scopesim/optics/optical_element.py +++ b/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): diff --git a/scopesim/optics/optics_manager.py b/scopesim/optics/optics_manager.py index 3f6b728a..cc5242a6 100644 --- a/scopesim/optics/optics_manager.py +++ b/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,75 +150,90 @@ 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) @@ -226,6 +241,7 @@ def surfaces_table(self): @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 diff --git a/scopesim/tests/tests_optics/test_OpticalElement.py b/scopesim/tests/tests_optics/test_OpticalElement.py index b6ec10b1..73bea078 100644 --- a/scopesim/tests/tests_optics/test_OpticalElement.py +++ b/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")