diff --git a/scopesim/effects/psfs/analytical.py b/scopesim/effects/psfs/analytical.py index dc39f5207..e40b37c99 100644 --- a/scopesim/effects/psfs/analytical.py +++ b/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) diff --git a/scopesim/effects/selector_wheel.py b/scopesim/effects/selector_wheel.py index 7ee03f242..41c9970e7 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__) @@ -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.") diff --git a/scopesim/effects/ter_curves.py b/scopesim/effects/ter_curves.py index 340af7eab..e48cae5d2 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 scipy import integrate from astropy import units as u @@ -1557,6 +1558,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) @@ -1564,6 +1566,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 = {} @@ -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: diff --git a/scopesim/optics/optical_train.py b/scopesim/optics/optical_train.py index 8ef81cd9c..e49e71c5f 100644 --- a/scopesim/optics/optical_train.py +++ b/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 diff --git a/scopesim/tests/tests_effects/test_MoffatPSF.py b/scopesim/tests/tests_effects/test_MoffatPSF.py index 1119df8dd..17e01a245 100644 --- a/scopesim/tests/tests_effects/test_MoffatPSF.py +++ b/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]) 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 000000000..ce308f392 --- /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