diff --git a/scopesim/effects/selector_wheel.py b/scopesim/effects/selector_wheel.py index 333635bb..c5869257 100644 --- a/scopesim/effects/selector_wheel.py +++ b/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.") diff --git a/scopesim/effects/ter_curves.py b/scopesim/effects/ter_curves.py index e35d2129..f4a6e421 100644 --- a/scopesim/effects/ter_curves.py +++ b/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,6 +1565,7 @@ 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) @@ -1571,6 +1573,9 @@ class DichroicTree(Effect): 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: diff --git a/scopesim/optics/optical_train.py b/scopesim/optics/optical_train.py index 9aa230d2..ebad53f7 100644 --- a/scopesim/optics/optical_train.py +++ b/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, diff --git a/scopesim/tests/tests_effects/test_MoffatPSF.py b/scopesim/tests/tests_effects/test_MoffatPSF.py index 6fe78db4..03b7bbcc 100644 --- a/scopesim/tests/tests_effects/test_MoffatPSF.py +++ b/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]) diff --git a/scopesim/tests/tests_effects/test_imaging_branches.py b/scopesim/tests/tests_effects/test_imaging_branches.py new file mode 100644 index 00000000..ce308f39 --- /dev/null +++ b/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 diff --git a/scopesim/tests/tests_optics/test_adc_slit_illumination.py b/scopesim/tests/tests_optics/test_adc_slit_illumination.py index 4dae1e4d..3df01129 100644 --- a/scopesim/tests/tests_optics/test_adc_slit_illumination.py +++ b/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)