Add zimager support with selector wheel and dichroic tree fixes.

scopesim/effects/psfs/analytical.py

@@ -258,19 +258,7 @@ def __init__(self, **kwargs):
         self._fwhm_interp = None
 
     def fwhm(self, wavelengths):
-        fwhm = self.meta["fwhm"]
-        if isinstance(fwhm, dict) and {"pivot_wave", "pivot_wave_unit"}.issubset(fwhm):
-            target, location, time, _ = get_observation_info_from_cmds(self.cmds)
-            key = tuple(fwhm.items()), target, location, time
-        elif isinstance(fwhm, dict):
-            key = tuple(fwhm.items())
-        else:
-            key = fwhm
-
-        if self._fwhm_interp is None or self._fwhm_interp[1] != key:
-            self._fwhm_interp = self.get_fwhm_interp(), key
-
-        return self._fwhm_interp[0](wavelengths)
+        return self.get_fwhm_interp()(wavelengths)
 
     @property
     def alpha(self):
@@ -282,33 +270,47 @@ def get_fwhm_interp(self):
         Overwrite this function for subclassing if FWHM input options change.
         """
         fwhm = self.meta["fwhm"]
+        if isinstance(fwhm, dict) and {"pivot_wave", "pivot_wave_unit"}.issubset(fwhm):
+            target, location, time, _ = get_observation_info_from_cmds(self.cmds)
+            zenith_angle = get_zenith_angle(target, location, time)
+            # The natural seeing law depends on the observation only through zenith angle.
+            key = tuple(fwhm.items()), zenith_angle
+        elif isinstance(fwhm, dict):
+            key = tuple(fwhm.items())
+        else:
+            key = fwhm
+
+        if self._fwhm_interp is not None and self._fwhm_interp[1] == key:
+            return self._fwhm_interp[0]
+
         if isinstance(fwhm, dict):
             logger.info("dict supplied for FWHM")
             if check_keys(fwhm, {"seeing", "seeing_unit", "pivot_wave", "pivot_wave_unit"}, action="warn"):
                 logger.info("seeing and pivot supplied, using natural scale seeing law")
-                target, location, time, _ = get_observation_info_from_cmds(self.cmds)
-                zenith_angle = get_zenith_angle(target, location, time)
                 seeing = (fwhm["seeing"] * u.Unit(fwhm["seeing_unit"])).to(u.arcsec)
                 pivot = (fwhm["pivot_wave"] * u.Unit(fwhm["pivot_wave_unit"])).to(u.um)
 
-                return partial(self.natural_scale, seeing=seeing, pivot=pivot, zenith_angle=zenith_angle*u.deg)
+                interp = partial(self.natural_scale, seeing=seeing, pivot=pivot, zenith_angle=zenith_angle*u.deg)
 
             elif check_keys(fwhm, {"seeing", "seeing_unit"}, action="error"):
                 logger.info("only seeing supplied, FWHM is wavelength independent")
                 seeing = (fwhm["seeing"] * u.Unit(fwhm["seeing_unit"])).to_value(u.arcsec)
-                return lambda wavelengths: np.full(wavelengths.shape, seeing, dtype=float) * u.arcsec
+                interp = lambda wavelengths: np.full(wavelengths.shape, seeing, dtype=float) * u.arcsec
 
-        if isinstance(fwhm, (int, float)):
+        elif isinstance(fwhm, (int, float)):
             logger.info("float value supplied, assuming arcsec")
-            return lambda wavelengths: np.full(wavelengths.shape, fwhm, dtype=float) * u.arcsec
+            interp = lambda wavelengths: np.full(wavelengths.shape, fwhm, dtype=float) * u.arcsec
 
-        if isinstance(fwhm, str):
+        elif isinstance(fwhm, str):
             logger.info("filename supplied for FWHM")
             self.table = DataContainer(filename=find_file(from_currsys(fwhm, self.cmds))).table
-            interp = self.fwhm_from_table(self.table)
-            return lambda wavelengths: interp(wavelengths) * u.arcsec
+            table_interp = self.fwhm_from_table(self.table)
+            interp = lambda wavelengths: table_interp(wavelengths) * u.arcsec
+        else:
+            raise TypeError("fwhm kwarg must be of type dict or float or str")
 
-        raise TypeError("fwhm kwarg must be of type dict or float or str")
+        self._fwhm_interp = interp, key
+        return interp
 
     def get_kernel(self, fov):
         pixel_scale = fov.header["CDELT1"] * u.deg.to(u.arcsec)

scopesim/effects/selector_wheel.py

@@ -19,6 +19,7 @@
 from ..optics.fov import FieldOfView
 from ..optics.image_plane import ImagePlane
 from ..detector.detector import Detector
+from ..source.source import Source
 
 logger = get_logger(__name__)
 
@@ -78,8 +79,14 @@ def __init__(self, **kwargs):
         self.z_order = self._resolve_z_order()
 
 
-    def apply_to(self, obj, **kwargs):
-        """Based on the selector_key's value in obj.meta, apply the corresponding effect from the wheel."""
+    def apply_to(self, obj, fov=None, **kwargs):
+        """Select from obj.meta, or fov.meta when applying branch effects to a Source."""
+        if isinstance(obj, Source) and fov is not None:
+            selector_value = fov.meta[self.meta["selector_key"]]
+            effect_to_apply = self.get_effect(selector_value)
+            if effect_to_apply is not None:
+                obj = effect_to_apply.apply_to(obj, fov=fov, **kwargs)
+
         if isinstance(obj, FieldOfView):
             if self.meta['selector_key'] not in obj.meta.keys():
                 raise ValueError(f"Selector key {self.meta['selector_key']} not found in FieldOfView meta.")

scopesim/effects/ter_curves.py

@@ -4,6 +4,7 @@
 import warnings
 from typing import ClassVar
 from collections.abc import Collection, Iterable
+from numbers import Integral
 import numpy as np
 from scipy import integrate
 from astropy import units as u
@@ -1557,13 +1558,17 @@ class DichroicTree(Effect):
     ```
 
     This effect creates a volume of FoVs, one for each row in the dichroic tree table.
+    For imaging, pass a Source with its fov context to apply the branch throughput before rendering.
     """
     required_keys = {"tree_filename", "dichroic_names", "filename_format"}
     z_order: ClassVar[tuple[int, ...]] = (117, 217, 617)
 
     def __init__(self, cmds=None, **kwargs):
         check_keys(kwargs, self.required_keys, action="error")
         super().__init__(cmds=cmds, **kwargs)
+        if "z_order" in self.meta:
+            z_order = self.meta["z_order"]
+            self.z_order = (int(z_order),) if isinstance(z_order, Integral) else tuple(z_order)
 
         # Load all the dichroic TER curves
         self.dichroics = {}
@@ -1580,7 +1585,7 @@ def __init__(self, cmds=None, **kwargs):
         tree = DataContainer(filename=tree_path)
         self.table = tree.table
 
-    def apply_to(self, obj, **kwargs):
+    def apply_to(self, obj, fov=None, **kwargs):
         """
         To the object, loop through each row of the table and apply the dichroics (defined by columns)
         with the corresponding action defined in the table cell.
@@ -1618,9 +1623,11 @@ def apply_to(self, obj, **kwargs):
             obj.volumes = new_vols
 
         # 2. During observe
-        if isinstance(obj, FieldOfView):
+        if isinstance(obj, FieldOfView) or (isinstance(obj, Source) and fov is not None):
             logger.debug("Executing %s, observe", self.meta['name'])
-            obj_id = obj.meta.get(id_colname, None)
+            branch_fov = obj if isinstance(obj, FieldOfView) else fov
+            obj_id = branch_fov.meta[id_colname]
+            setattr(branch_fov, id_colname, obj_id)
             row = self.table[self.table[id_colname] == obj_id]
 
             if len(row) != 1:

scopesim/optics/optical_train.py

@@ -23,7 +23,7 @@
 from .image_plane import ImagePlane
 from ..commands.user_commands import UserCommands
 from ..detector import DetectorManager
-from ..effects import ExtraFitsKeywords
+from ..effects import DichroicTree, ExtraFitsKeywords, SelectorWheel
 from ..utils import from_currsys, top_level_catch, get_logger
 from ..source.source_templates import empty_sky
 from ..source.source_fields import SpectrumSourceField
@@ -37,7 +37,12 @@
 N_PROCESSES = mp.cpu_count() - 1
 USE_MULTIPROCESSING = False
 
-def extract_source(fov, source):
+def extract_source(fov, source, source_effects=()):
+    # Branch spectra must be filtered before extraction clips them to the FOV wavelength range.
+    if source_effects:
+        source = source.make_copy()
+        for effect in source_effects:
+            source = effect.apply_to(source, fov=fov)
     fov.extract_from(source)
     return fov
 
@@ -538,7 +543,16 @@ def observe(self, orig_source=None, update=True, **kwargs):
         source = self.prepare_source(source)
 
         # [1D - transmission curves]
-        for effect in self.optics_manager.source_effects:
+        source_effects = self.optics_manager.source_effects
+        branch_effects = []
+        if not self.fov_manager.is_spectroscope:
+            for ii, effect in enumerate(source_effects):
+                if isinstance(effect, (DichroicTree, SelectorWheel)):
+                    # Keep downstream attenuation and emission in optical order for each branch.
+                    source_effects, branch_effects = source_effects[:ii], source_effects[ii:]
+                    break
+
+        for effect in source_effects:
             source = effect.apply_to(source)
 
         # [3D - Atmospheric shifts, PSF, NCPAs, Grating shift/distortion]
@@ -552,7 +566,7 @@ def observe(self, orig_source=None, update=True, **kwargs):
 
             with mp.Pool(processes=N_PROCESSES) as pool:
                 fovs = pool.starmap(extract_source,
-                                    zip(fovs, [source] * len(fovs)))
+                                    zip(fovs, [source] * len(fovs), [branch_effects] * len(fovs)))
 
             with mp.Pool(processes=N_PROCESSES) as pool:
                 fovs = pool.starmap(view_fov,
@@ -568,7 +582,7 @@ def observe(self, orig_source=None, update=True, **kwargs):
             fovs = self.fov_manager.fovs
             nobar = len(fovs) <= 1
             for fov in tqdm(fovs, desc=" FOVs", position=0, disable=nobar):
-                fov.extract_from(source)
+                fov = extract_source(fov, source, branch_effects)
 
                 fov.view()
                 foveffs = self.optics_manager.fov_effects

scopesim/tests/tests_effects/test_MoffatPSF.py

@@ -1,10 +1,14 @@
 import logging
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
 from astropy import units as u
+from astropy.coordinates import EarthLocation, SkyCoord
+from astropy.table import Table
+from astropy.time import Time
 
-from scopesim.effects.psfs.analytical import MoffatPSF
+from scopesim.effects.psfs.analytical import AOEnhanceablePSF, MoffatPSF
 
 
 class SimpleFov:
@@ -87,3 +91,104 @@ def test_legacy_renormalization_can_be_requested():
     kernel = psf.get_kernel(SimpleFov())
 
     assert np.sum(kernel) == pytest.approx(1.0)
+
+
+@pytest.fixture
+def seeing_law():
+    return {"seeing": 0.7, "seeing_unit": "arcsec", "pivot_wave": 500, "pivot_wave_unit": "nm"}
+
+
+@pytest.fixture
+def observing_conditions(monkeypatch):
+    location = EarthLocation.from_geodetic(-155 * u.deg, 20 * u.deg, 4000 * u.m)
+    conditions = {"time": Time("2026-01-01T00:00:00"), "alt": 60 * u.deg}
+
+    def observation_info(cmds):
+        time = conditions["time"]
+        target = SkyCoord(alt=conditions["alt"], az=0 * u.deg, frame="altaz", location=location, obstime=time)
+        return target, location, time, None
+
+    monkeypatch.setattr("scopesim.effects.psfs.analytical.get_observation_info_from_cmds", observation_info)
+    return conditions
+
+
+@pytest.mark.parametrize("psf_class", [MoffatPSF, AOEnhanceablePSF])
+def test_observation_configuration_is_only_required_when_evaluating(psf_class, seeing_law, monkeypatch):
+    observation_info = Mock(side_effect=KeyError("!OBS.airmass"))
+    monkeypatch.setattr("scopesim.effects.psfs.analytical.get_observation_info_from_cmds", observation_info)
+    kwargs = {"alpha": 4.765, "fwhm": seeing_law}
+    if psf_class is AOEnhanceablePSF:
+        kwargs.update(ao_table="unused.dat", enable_ao=False, is_absolute=False)
+    psf = psf_class(**kwargs)
+    observation_info.assert_not_called()
+    with pytest.raises(KeyError, match="airmass"):
+        psf.fwhm(np.array([0.5]) * u.um)
+
+
+@pytest.mark.parametrize("fwhm", [0.7, {"seeing": 700, "seeing_unit": "mas"}])
+def test_constant_seeing_is_evaluated_without_observation_commands(fwhm, monkeypatch):
+    observation_info = Mock(side_effect=KeyError("!OBS.airmass"))
+    monkeypatch.setattr("scopesim.effects.psfs.analytical.get_observation_info_from_cmds", observation_info)
+    psf = MoffatPSF(alpha=4.765, fwhm=fwhm)
+    np.testing.assert_allclose(psf.fwhm(np.array([0.5, 1.0]) * u.um).to_value(u.arcsec), [0.7, 0.7])
+    observation_info.assert_not_called()
+
+
+def test_changing_observation_times_reuses_same_seeing_law(seeing_law, observing_conditions):
+    psf = MoffatPSF(alpha=4.765, fwhm=seeing_law)
+    wavelengths = np.array([0.5, 1.0]) * u.um
+    first = psf.fwhm(wavelengths)
+    cached_interp = psf._fwhm_interp[0]
+    observing_conditions["time"] = observing_conditions["time"] + 1 * u.hour
+    second = psf.fwhm(wavelengths)
+    np.testing.assert_allclose(second.to_value(u.arcsec), first.to_value(u.arcsec))
+    assert psf._fwhm_interp[0] is cached_interp
+
+
+def test_zenith_angle_change_updates_seeing_law(seeing_law, observing_conditions):
+    psf = MoffatPSF(alpha=4.765, fwhm=seeing_law)
+    wavelengths = np.array([0.5]) * u.um
+    first = psf.fwhm(wavelengths)
+    cached_interp = psf._fwhm_interp[0]
+    observing_conditions["alt"] = 30 * u.deg
+    second = psf.fwhm(wavelengths)
+    expected_ratio = (np.cos(30 * u.deg) / np.cos(60 * u.deg))**0.6
+    np.testing.assert_allclose(second.to_value(u.arcsec) / first.to_value(u.arcsec), expected_ratio)
+    assert psf._fwhm_interp[0] is not cached_interp
+
+
+@pytest.mark.parametrize("parameter, value, expected_ratio", [("seeing", 0.9, 0.9 / 0.7), ("pivot_wave", 1000, 2**0.2)])
+def test_seeing_parameters_update_cached_law(parameter, value, expected_ratio, seeing_law, observing_conditions):
+    psf = MoffatPSF(alpha=4.765, fwhm=seeing_law)
+    wavelengths = np.array([1.0]) * u.um
+    first = psf.fwhm(wavelengths)
+    psf.meta["fwhm"][parameter] = value
+    second = psf.fwhm(wavelengths)
+    np.testing.assert_allclose(second.to_value(u.arcsec) / first.to_value(u.arcsec), expected_ratio)
+
+
+def test_tabulated_fwhm_without_observation_commands(tmp_path, monkeypatch):
+    observation_info = Mock(side_effect=KeyError("!OBS.airmass"))
+    monkeypatch.setattr("scopesim.effects.psfs.analytical.get_observation_info_from_cmds", observation_info)
+    filename = tmp_path / "fwhm.ecsv"
+    table = Table({"wavelength": np.array([0.5, 1.0, 1.5, 2.0]) * u.um, "fwhm": [0.5, 0.6, 0.7, 0.8]})
+    table.write(filename, format="ascii.ecsv")
+    psf = MoffatPSF(alpha=4.765, fwhm=str(filename))
+    np.testing.assert_allclose(psf.fwhm(np.array([0.5, 1.0]) * u.um).to_value(u.arcsec), [0.5, 0.6])
+    observation_info.assert_not_called()
+
+
+def test_ao_toggle_preserves_natural_and_absolute_widths(tmp_path):
+    filename = tmp_path / "ao.ecsv"
+    table = Table({"wavelength": np.array([0.5, 1.0, 1.5, 2.0]) * u.um, "fwhm": [0.5, 0.5, 0.5, 0.5]})
+    table.write(filename, format="ascii.ecsv")
+    psf = AOEnhanceablePSF(alpha=4.765, fwhm=0.7, ao_table=str(filename), enable_ao=False, is_absolute=False)
+    wavelengths = np.array([1.0]) * u.um
+    np.testing.assert_allclose(psf.fwhm(wavelengths).to_value(u.arcsec), [0.7])
+    psf.meta["enable_ao"] = True
+    np.testing.assert_allclose(psf.fwhm(wavelengths).to_value(u.arcsec), [0.35])
+    psf.meta["is_absolute"] = True
+    np.testing.assert_allclose(psf.fwhm(wavelengths).to_value(u.arcsec), [0.5])
+    psf.meta["fwhm"] = 0.9
+    psf.meta["enable_ao"] = False
+    np.testing.assert_allclose(psf.fwhm(wavelengths).to_value(u.arcsec), [0.9])