[27/29] Apply imaging branch throughput before wavelength integration

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__)
 
@@ -80,8 +81,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 numpy.typing import ArrayLike, NDArray
 from scipy import integrate
@@ -1564,13 +1565,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 = {}
@@ -1587,7 +1592,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.
@@ -1625,9 +1630,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

@@ -38,8 +38,14 @@
 N_PROCESSES = mp.cpu_count() - 1
 USE_MULTIPROCESSING = False
 
-def extract_source(fov, source):
-    fov.extract_from(source)
+def extract_source(fov, source, source_effects=(), spatial_header=None, _spectral_wavesets=None):
+    if source_effects:
+        source = source.make_copy()
+        for effect in source_effects:
+            source = effect.apply_to(source, fov=fov)
+        # Branch filtering changes the spectrum and its native wavelength grid.
+        _spectral_wavesets = None
+    fov.extract_from(source, spatial_header=spatial_header, _spectral_wavesets=_spectral_wavesets)
     return fov
 
 def view_fov(fov, hdu_type):
@@ -535,7 +541,15 @@ 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)):
+                    # Attenuation and emission downstream of a split belong to each branch.
+                    source_effects, branch_effects = source_effects[:ii], source_effects[ii:]
+                    break
+        for effect in source_effects:
             source = effect.apply_to(source)
 
         # This source is a private copy, after all source effects. FOV extraction
@@ -558,7 +572,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,
@@ -614,9 +628,9 @@ def observe(self, orig_source=None, update=True, **kwargs):
                     spatial_header = foveffs[psf_index].padded_header(fov, point_sources=point_sources)
                     for effect in spatial_shifts:
                         spatial_header = effect._padded_header(spatial_header, effect.get_shifts(fov))
-                    fov.extract_from(source, spatial_header=spatial_header, _spectral_wavesets=spectral_wavesets)
+                    extract_source(fov, source, branch_effects, spatial_header=spatial_header, _spectral_wavesets=spectral_wavesets)
                 else:
-                    fov.extract_from(source, _spectral_wavesets=spectral_wavesets)
+                    extract_source(fov, source, branch_effects, _spectral_wavesets=spectral_wavesets)
                     fov.view()
                 nobar = len(foveffs) <= 1
                 for effect in tqdm(foveffs, disable=nobar,

scopesim/tests/tests_effects/test_MoffatPSF.py

@@ -1,4 +1,5 @@
 import logging
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
@@ -251,3 +252,33 @@ def test_ao_cache_tracks_resolved_path_contents_and_alpha(tmp_path):
     np.testing.assert_allclose(psf.fwhm(wave).value, .2)
     psf.meta['alpha'] = 3.
     assert psf.alpha == 3.
+
+
+@pytest.mark.parametrize("psf_class", [MoffatPSF, AOEnhanceablePSF])
+def test_observation_configuration_is_only_required_when_evaluating(psf_class, 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": 0.7, "seeing_unit": "arcsec", "pivot_wave": 500, "pivot_wave_unit": "nm"}}
+    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)
+
+
+
+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])

scopesim/tests/tests_effects/test_imaging_branches.py

@@ -0,0 +1,115 @@
+"""Branch throughput must be applied before integrating an imaging spectrum."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from scopesim.commands import UserCommands
+from scopesim.effects import DichroicTree, SelectorWheel, TERCurve
+from scopesim.optics.image_plane import ImagePlane
+from scopesim.optics.optical_train import OpticalTrain, extract_source
+from scopesim.source.source import Source
+from scopesim.tests.mocks.py_objects.fov_objects import _centre_fov
+from scopesim.tests.mocks.py_objects.source_objects import _table_source
+
+
+@pytest.fixture
+def tree(tmp_path):
+    (tmp_path / 'tree.dat').write_text('image_plane_id first second\n0 R X\n1 T R\n2 T T\n')
+    for name, transmission in [('first', 0.75), ('second', 0.6)]:
+        curve = '# wavelength_unit: um\nwavelength transmission reflection emissivity\n'
+        curve += f'0.5 {transmission} {1 - transmission} 0\n2.5 {transmission} {1 - transmission} 0\n'
+        (tmp_path / f'{name}.dat').write_text(curve)
+    return DichroicTree(name='tree', tree_filename=str(tmp_path / 'tree.dat'), dichroic_names=['first', 'second'], filename_format=str(tmp_path / '{}.dat'))
+
+
+@pytest.fixture
+def selector():
+    wheel = []
+    for arm, transmission in enumerate([0.25, 0.75]):
+        curve = {'wavelength': [0.5, 2.5], 'transmission': [transmission, transmission], 'emissivity': [0, 0]}
+        wheel.append({'selector_value': arm, 'effect_class': 'QuantumEfficiencyCurve', 'effect_kwargs': {'array_dict': curve, 'wavelength_unit': 'um'}})
+    return SelectorWheel(selector_key='image_plane_id', wheel=wheel, allowed_missing_selector_values=[2])
+
+
+@pytest.fixture
+def source():
+    source = _table_source()
+    background = TERCurve(array_dict={'wavelength': [0.5, 2.5], 'transmission': [1, 1], 'emission': [1e-8, 1e-8]}, wavelength_unit='um', emission_unit='photlam')
+    source.append(background.background_source[0])
+    return source
+
+
+@pytest.mark.parametrize('effect_name, transmissions', [('tree', [0.25, 0.3, 0.45]), ('selector', [0.25, 0.75, 1.0])])
+def test_branch_attenuates_coloured_sources_and_background_independently(request, effect_name, transmissions, source):
+    effect = request.getfixturevalue(effect_name)
+    reference = extract_source(_centre_fov(n=10), source).make_hdu().data
+    for arm, transmission in enumerate(transmissions):
+        fov = _centre_fov(n=10)
+        fov.meta['image_plane_id'] = arm
+        extract_source(fov, source, [effect])
+        np.testing.assert_allclose(fov.make_hdu().data, transmission * reference)
+    # Rendering each branch must leave the shared source, including its sky, intact.
+    np.testing.assert_array_equal(extract_source(_centre_fov(n=10), source).make_hdu().data, reference)
+
+
+@pytest.mark.parametrize('effect_name', ['tree', 'selector'])
+def test_source_without_fov_context_is_unchanged(request, effect_name, source):
+    effect = request.getfixturevalue(effect_name)
+    reference = extract_source(_centre_fov(n=10), source).make_hdu().data
+    effect.apply_to(source)
+    np.testing.assert_array_equal(extract_source(_centre_fov(n=10), source).make_hdu().data, reference)
+
+
+@pytest.mark.parametrize('effect_name', ['tree', 'selector'])
+def test_missing_branch_id_fails(request, effect_name, source):
+    effect = request.getfixturevalue(effect_name)
+    fov = _centre_fov(n=10)
+    fov.meta.clear()
+    with pytest.raises(KeyError, match='image_plane_id'):
+        effect.apply_to(source, fov=fov)
+
+
+def test_observe_keeps_downstream_emission_after_branch_split(tree, source):
+    common = TERCurve(array_dict={'wavelength': [0.5, 2.5], 'transmission': [0.8, 0.8], 'emissivity': [0, 0]}, wavelength_unit='um')
+    downstream = TERCurve(array_dict={'wavelength': [0.5, 2.5], 'transmission': [0.5, 0.5], 'emission': [2e-8, 2e-8]}, wavelength_unit='um', emission_unit='photlam')
+    train = OpticalTrain()
+    train.cmds = UserCommands()
+    fovs = [_centre_fov(n=10) for arm in range(3)]
+    for arm, fov in enumerate(fovs):
+        fov.meta['image_plane_id'] = arm
+    train.fov_manager = SimpleNamespace(fovs=fovs, is_spectroscope=False)
+    train.optics_manager = SimpleNamespace(source_effects=[common, tree, downstream], fov_effects=[], image_plane_effects=[])
+    train.image_planes = [ImagePlane(fov.header) for fov in fovs]
+
+    prepared = train.prepare_source(source.make_copy())
+    reference = extract_source(_centre_fov(n=10), prepared).make_hdu().data
+    emission = extract_source(_centre_fov(n=10), downstream.background_source[0]).make_hdu().data
+    train.observe(source, update=False)
+    for arm, transmission in enumerate([0.25, 0.3, 0.45]):
+        expected = reference * 0.8 * transmission * 0.5 + emission
+        np.testing.assert_allclose(train._last_fovs[arm].hdu.data, expected)
+        assert train._last_fovs[arm].image_plane_id == arm
+
+
+@pytest.mark.parametrize('spectrum_id, expected', [(0, 0.5), (1, 0.5 + 1 / 24), (2, 0.5 - 1 / 24)])
+def test_wavelength_dependent_branch_preserves_source_colour(spectrum_id, expected):
+    spectrum = _table_source().fields[0].spectra[spectrum_id]
+    source = Source(x=[0], y=[0], ref=[0], weight=[1], spectra=[spectrum])
+    curve = {'wavelength': [0.5, 2.5], 'transmission': [0, 1], 'emissivity': [0, 0]}
+    selector = SelectorWheel(selector_key='image_plane_id', wheel=[{'selector_value': 0, 'effect_class': 'QuantumEfficiencyCurve', 'effect_kwargs': {'array_dict': curve, 'wavelength_unit': 'um'}}])
+    reference = extract_source(_centre_fov(n=10), source).make_hdu().data.sum()
+    fov = _centre_fov(n=10)
+    fov.meta['image_plane_id'] = 0
+    measured = extract_source(fov, source, [selector]).make_hdu().data.sum() / reference
+    # Linear spectra and throughput give these analytic integrals over 1--2 um.
+    # The input spectra are sampled every 0.02 um.
+    np.testing.assert_allclose(measured, expected, rtol=0, atol=1e-4)
+
+
+@pytest.mark.parametrize('z_order, expected', [([117, 217, 517], (117, 217, 517)), (517, (517,))])
+def test_dichroic_source_stage_is_opt_in(tree, z_order, expected):
+    assert tree.z_order == (117, 217, 617)
+    configured = DichroicTree(z_order=z_order, **{key: tree.meta[key] for key in tree.required_keys})
+    assert configured.z_order == expected

scopesim/tests/tests_optics/test_adc_slit_illumination.py

@@ -157,7 +157,7 @@ def test_observe_applies_shift_before_slit_through_wavelength_throughput(interve
             effects.append(Displacement(0., displacement, method='roll'))
         train = OpticalTrain()
         train.cmds = fov.cmds
-        train.fov_manager = SimpleNamespace(fovs=[fov])
+        train.fov_manager = SimpleNamespace(fovs=[fov], is_spectroscope=True)
         train.optics_manager = SimpleNamespace(source_effects=[], fov_effects=effects, image_plane_effects=[])
         train.image_planes = [SimpleNamespace(add=lambda hdu, **kwargs: images.append(hdu.data.copy()))]
         train.observe(source, update=False)