From 67b3545ed1c69b8c391078dcbe1f47090afc2ed4 Mon Sep 17 00:00:00 2001 From: Tim Beccue Date: Tue, 15 Sep 2026 16:21:25 -0600 Subject: [PATCH 1/4] Avoid rebuilding aligned single-component mosaics --- banzai/mosaic.py | 36 +++++++++++-- banzai/tests/test_mosaic_creator.py | 82 +++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 5 deletions(-) diff --git a/banzai/mosaic.py b/banzai/mosaic.py index 253ed7e9d..c9e69553c 100644 --- a/banzai/mosaic.py +++ b/banzai/mosaic.py @@ -14,13 +14,28 @@ def __init__(self, runtime_context): def do_stage(self, image): logger.info('Mosaicing image', image=image) + ccd_hdus = image.ccd_hdus mosaiced_detector_region = self.get_mosaic_detector_region(image) binned_shape = [length // binning for length, binning in zip(mosaiced_detector_region.shape, image.binning)] - mosaiced_data = CCDData(data=np.zeros(binned_shape, dtype=image.data_type), - meta=image.primary_hdu.meta) + mosaiced_data_section = Section(x_start=1, y_start=1, x_stop=binned_shape[1], y_stop=binned_shape[0]) + data_type = image.data_type + reuse_component = len(ccd_hdus) == 1 and self._can_reuse_component( + ccd_hdus[0], mosaiced_detector_region, mosaiced_data_section, data_type) + if reuse_component: + # Borrow the existing arrays while keeping the same primary-header normalization below. + component = ccd_hdus[0] + mosaiced_data = CCDData(data=component.data, meta=image.primary_hdu.meta, mask=component.mask, + uncertainty=component.uncertainty, memmap=False) + # Subsequent stages retain the usual policy of allocating new memory maps. + mosaiced_data.memmap = True + if component.uncertainty.dtype != data_type: + mosaiced_data.uncertainty = component.uncertainty.astype(data_type) + else: + mosaiced_data = CCDData(data=np.zeros(binned_shape, dtype=data_type), + meta=image.primary_hdu.meta) mosaiced_data.binning = image.binning mosaiced_data.detector_section = mosaiced_detector_region - mosaiced_data.data_section = Section(x_start=1, y_start=1, x_stop=binned_shape[1], y_stop=binned_shape[0]) + mosaiced_data.data_section = mosaiced_data_section mosaiced_data.name = 'SCI' mosaiced_data.gain = np.mean([data.meta.get('GAIN') for data in image.ccd_hdus]) @@ -31,8 +46,9 @@ def do_stage(self, image): mosaiced_data.meta['L1STATOV'] = '1' if any([data.meta.get('L1STATOV', '0') == '1' for data in image.ccd_hdus]) \ else '0', 'Status flag for overscan correction' - for i, data in enumerate(image.ccd_hdus): - mosaiced_data.copy_in(data) + for i, data in enumerate(ccd_hdus): + if not reuse_component: + mosaiced_data.copy_in(data) mosaiced_data.meta[f'OVERSCN{i + 1}'] = '{:0.2f}'.format(data.meta.get('OVERSCAN', 0.0)), \ 'Overscan value that was subtracted' image.remove(data) @@ -40,6 +56,16 @@ def do_stage(self, image): image.primary_hdu = mosaiced_data return image + @staticmethod + def _can_reuse_component(data, detector_section, data_section, data_type): + # A single component can still need cropping, flipping, or science/mask dtype conversion. + return (data.shape == data_section.shape + and data.data_section is not None + and data.data_section.to_region_keyword() == data_section.to_region_keyword() + and data.detector_section.to_region_keyword() == detector_section.to_region_keyword() + and data.dtype == data_type + and data.mask.dtype == np.uint8) + @staticmethod def get_mosaic_detector_region(image): x_detector_sections = [] diff --git a/banzai/tests/test_mosaic_creator.py b/banzai/tests/test_mosaic_creator.py index cc5ddc007..7916b15a4 100644 --- a/banzai/tests/test_mosaic_creator.py +++ b/banzai/tests/test_mosaic_creator.py @@ -1,6 +1,12 @@ +import tempfile +from unittest.mock import patch + import pytest import numpy as np +from astropy.io.fits import Header +from banzai.data import CCDData, HeaderOnly +from banzai.lco import LCOObservationFrame from banzai.utils.image_utils import Section from banzai.mosaic import MosaicCreator from banzai.tests.utils import FakeLCOObservationFrame, FakeCCDData @@ -45,6 +51,82 @@ def test_get_mosaic_detector_region(): assert MosaicCreator.get_mosaic_detector_region(image).shape == (2048, 2048) +def make_single_component_frame(detsec='[1:6,1:4]', datasec='[1:6,1:4]', binning='1 1', + separate_primary=False, uncertainty_dtype=np.float64): + meta = Header({'OBSTYPE': 'SUB_EXP', 'DAY-OBS': '20260915', 'MOLUID': 123, 'MOLFRNUM': 2, + 'FRMTOTAL': 10, 'GAIN': 1.0, 'SATURATE': 40000.0, 'MAXLIN': 30000.0, + 'RDNOISE': 3.0, 'CCDSUM': binning, 'DATASEC': datasec, 'DETSEC': detsec, + 'TRIMSEC': '[1:6,1:4]', 'OVERSCAN': 3.25, 'L1STATOV': '1'}) + pixels = np.arange(24, dtype=np.float64).reshape(4, 6) + component = CCDData(data=pixels, meta=meta, mask=pixels.astype(np.uint8) % 16, + uncertainty=(pixels / 10 + 1).astype(uncertainty_dtype), name='RAW') + hdus = [component] + if separate_primary: + primary_meta = meta.copy() + primary_meta.update({'GAIN': 2.0, 'SATURATE': 50000.0, 'MAXLIN': 45000.0, + 'L1STATOV': '0', 'OVERSCAN': 0.0, 'PRIMARY': 'preserve'}) + component.meta['AMPONLY'] = 'omit' + hdus.insert(0, HeaderOnly(meta=primary_meta, name='PRIMARY')) + return LCOObservationFrame(hdu_list=hdus, file_path='/tmp/single-component.fits') + + +@pytest.mark.parametrize(('separate_primary', 'detsec', 'binning', 'uncertainty_dtype'), [ + (False, '[1:6,1:4]', '1 1', np.float64), + (True, '[1:6,1:4]', '1 1', np.float64), + (True, '[101:112,201:208]', '2 2', np.float64), + (True, '[1:6,1:4]', '1 1', np.float32), +]) +def test_single_component_reuses_arrays_and_matches_general_mosaic( + monkeypatch, separate_primary, detsec, binning, uncertainty_dtype): + kwargs = dict(separate_primary=separate_primary, detsec=detsec, binning=binning, + uncertainty_dtype=uncertainty_dtype) + with monkeypatch.context() as general_path: + general_path.setattr(MosaicCreator, '_can_reuse_component', staticmethod(lambda *args: False)) + expected = MosaicCreator(None).do_stage(make_single_component_frame(**kwargs)) + + image = make_single_component_frame(**kwargs) + component = image.ccd_hdus[0] + original_meta = component.meta.copy() + with patch('banzai.data.tempfile.NamedTemporaryFile', wraps=tempfile.NamedTemporaryFile) as new_file: + actual = MosaicCreator(None).do_stage(image) + + assert actual is image + assert len(actual.ccd_hdus) == 1 + assert actual.primary_hdu.name == 'SCI' + assert actual.meta == expected.meta + assert component.meta == original_meta + assert component.memmap is True + assert actual.primary_hdu.memmap is True + assert actual.data is component.data + assert actual.mask is component.mask + if uncertainty_dtype == np.float64: + assert actual.uncertainty is component.uncertainty + assert new_file.call_count == 0 + else: + assert actual.uncertainty is not component.uncertainty + assert new_file.call_count == 1 + for name in ('data', 'mask', 'uncertainty'): + result_array = getattr(actual.primary_hdu, name) + expected_array = getattr(expected.primary_hdu, name) + assert result_array.dtype == expected_array.dtype + np.testing.assert_array_equal(result_array, expected_array) + + +@pytest.mark.parametrize(('detsec', 'datasec', 'expected_slice'), [ + ('[6:1,1:4]', '[1:6,1:4]', (slice(None), slice(None, None, -1))), + ('[1:6,1:4]', '[6:1,1:4]', (slice(None), slice(None, None, -1))), + ('[1:4,1:4]', '[2:5,1:4]', (slice(None), slice(1, 5))), +]) +def test_single_component_still_crops_and_flips(detsec, datasec, expected_slice): + image = make_single_component_frame(detsec=detsec, datasec=datasec) + component = image.ccd_hdus[0] + actual = MosaicCreator(None).do_stage(image) + + assert actual.data is not component.data + for name in ('data', 'mask', 'uncertainty'): + np.testing.assert_array_equal(getattr(actual.primary_hdu, name), getattr(component, name)[expected_slice]) + + def test_mosaic_maker(set_random_seed): detsecs = [['[1:100,1:100]', '[1:100,200:101]', '[200:101,1:100]', '[200:101,200:101]'], ['[1:200,400:201]', '[1:200,1:200]', '[400:201,400:201]', '[400:201,1:200]'], From d997a4869a1787585f3904217c8e8aab9cb5e4e7 Mon Sep 17 00:00:00 2001 From: Tim Beccue Date: Tue, 15 Sep 2026 16:43:37 -0600 Subject: [PATCH 2/4] Restrict mosaic reuse to compatible arrays --- banzai/mosaic.py | 8 ++++++-- banzai/tests/test_mosaic_creator.py | 32 +++++++++++++++++++++++++++-- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/banzai/mosaic.py b/banzai/mosaic.py index c9e69553c..3f76d8a0b 100644 --- a/banzai/mosaic.py +++ b/banzai/mosaic.py @@ -59,12 +59,16 @@ def do_stage(self, image): @staticmethod def _can_reuse_component(data, detector_section, data_section, data_type): # A single component can still need cropping, flipping, or science/mask dtype conversion. - return (data.shape == data_section.shape + # Other science dtypes can be promoted when the normal constructor initializes uncertainties. + return (data_type == np.float64 + and data.shape == data_section.shape and data.data_section is not None and data.data_section.to_region_keyword() == data_section.to_region_keyword() and data.detector_section.to_region_keyword() == detector_section.to_region_keyword() and data.dtype == data_type - and data.mask.dtype == np.uint8) + and data.mask.dtype == np.uint8 + and all(array.flags.c_contiguous and array.flags.writeable + for array in (data.data, data.mask, data.uncertainty))) @staticmethod def get_mosaic_detector_region(image): diff --git a/banzai/tests/test_mosaic_creator.py b/banzai/tests/test_mosaic_creator.py index 7916b15a4..2a6f304dd 100644 --- a/banzai/tests/test_mosaic_creator.py +++ b/banzai/tests/test_mosaic_creator.py @@ -52,12 +52,12 @@ def test_get_mosaic_detector_region(): def make_single_component_frame(detsec='[1:6,1:4]', datasec='[1:6,1:4]', binning='1 1', - separate_primary=False, uncertainty_dtype=np.float64): + separate_primary=False, uncertainty_dtype=np.float64, science_dtype=np.float64): meta = Header({'OBSTYPE': 'SUB_EXP', 'DAY-OBS': '20260915', 'MOLUID': 123, 'MOLFRNUM': 2, 'FRMTOTAL': 10, 'GAIN': 1.0, 'SATURATE': 40000.0, 'MAXLIN': 30000.0, 'RDNOISE': 3.0, 'CCDSUM': binning, 'DATASEC': datasec, 'DETSEC': detsec, 'TRIMSEC': '[1:6,1:4]', 'OVERSCAN': 3.25, 'L1STATOV': '1'}) - pixels = np.arange(24, dtype=np.float64).reshape(4, 6) + pixels = np.arange(24, dtype=science_dtype).reshape(4, 6) component = CCDData(data=pixels, meta=meta, mask=pixels.astype(np.uint8) % 16, uncertainty=(pixels / 10 + 1).astype(uncertainty_dtype), name='RAW') hdus = [component] @@ -127,6 +127,34 @@ def test_single_component_still_crops_and_flips(detsec, datasec, expected_slice) np.testing.assert_array_equal(getattr(actual.primary_hdu, name), getattr(component, name)[expected_slice]) +@pytest.mark.parametrize('science_dtype', [np.float32, np.int32]) +def test_single_component_preserves_promoted_uncertainty_dtype(science_dtype): + image = make_single_component_frame(science_dtype=science_dtype) + component = image.ccd_hdus[0] + actual = MosaicCreator(None).do_stage(image) + + assert actual.uncertainty.dtype == np.float64 + np.testing.assert_array_equal(actual.data, component.data) + np.testing.assert_array_equal(actual.uncertainty, component.uncertainty) + + +@pytest.mark.parametrize('storage', ['read_only', 'strided']) +def test_single_component_preserves_writable_contiguous_output(storage): + image = make_single_component_frame() + component = image.ccd_hdus[0] + if storage == 'read_only': + component.data.flags.writeable = False + else: + component.data = component.data[:, ::-1] + actual = MosaicCreator(None).do_stage(image) + + for name in ('data', 'mask', 'uncertainty'): + array = getattr(actual.primary_hdu, name) + assert array.flags.writeable + assert array.flags.c_contiguous + np.testing.assert_array_equal(array, getattr(component, name)) + + def test_mosaic_maker(set_random_seed): detsecs = [['[1:100,1:100]', '[1:100,200:101]', '[200:101,1:100]', '[200:101,200:101]'], ['[1:200,400:201]', '[1:200,1:200]', '[400:201,400:201]', '[400:201,1:200]'], From ae971687659bc062cacbf1d9c1b541d2d1e8b61d Mon Sep 17 00:00:00 2001 From: Tim Beccue Date: Tue, 15 Sep 2026 17:04:26 -0600 Subject: [PATCH 3/4] Explain single-component mosaic array reuse --- banzai/mosaic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/banzai/mosaic.py b/banzai/mosaic.py index 3f76d8a0b..0efe9a908 100644 --- a/banzai/mosaic.py +++ b/banzai/mosaic.py @@ -21,6 +21,7 @@ def do_stage(self, image): data_type = image.data_type reuse_component = len(ccd_hdus) == 1 and self._can_reuse_component( ccd_hdus[0], mosaiced_detector_region, mosaiced_data_section, data_type) + # Save time by reusing compatible arrays when a single component already matches the output mosaic layout. if reuse_component: # Borrow the existing arrays while keeping the same primary-header normalization below. component = ccd_hdus[0] From 688b322e9a5dd80e0d51cd49fabfad6140b98cd8 Mon Sep 17 00:00:00 2001 From: Tim Beccue Date: Tue, 29 Sep 2026 15:12:05 -0600 Subject: [PATCH 4/4] Move single-component check into _can_reuse_component; add y-flip tests --- banzai/mosaic.py | 9 ++++++--- banzai/tests/test_mosaic_creator.py | 2 ++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/banzai/mosaic.py b/banzai/mosaic.py index 0efe9a908..5a12ad0e9 100644 --- a/banzai/mosaic.py +++ b/banzai/mosaic.py @@ -19,8 +19,8 @@ def do_stage(self, image): binned_shape = [length // binning for length, binning in zip(mosaiced_detector_region.shape, image.binning)] mosaiced_data_section = Section(x_start=1, y_start=1, x_stop=binned_shape[1], y_stop=binned_shape[0]) data_type = image.data_type - reuse_component = len(ccd_hdus) == 1 and self._can_reuse_component( - ccd_hdus[0], mosaiced_detector_region, mosaiced_data_section, data_type) + reuse_component = self._can_reuse_component(ccd_hdus, mosaiced_detector_region, mosaiced_data_section, + data_type) # Save time by reusing compatible arrays when a single component already matches the output mosaic layout. if reuse_component: # Borrow the existing arrays while keeping the same primary-header normalization below. @@ -58,7 +58,10 @@ def do_stage(self, image): return image @staticmethod - def _can_reuse_component(data, detector_section, data_section, data_type): + def _can_reuse_component(ccd_hdus, detector_section, data_section, data_type): + if len(ccd_hdus) != 1: + return False + data = ccd_hdus[0] # A single component can still need cropping, flipping, or science/mask dtype conversion. # Other science dtypes can be promoted when the normal constructor initializes uncertainties. return (data_type == np.float64 diff --git a/banzai/tests/test_mosaic_creator.py b/banzai/tests/test_mosaic_creator.py index 2a6f304dd..990a539df 100644 --- a/banzai/tests/test_mosaic_creator.py +++ b/banzai/tests/test_mosaic_creator.py @@ -115,6 +115,8 @@ def test_single_component_reuses_arrays_and_matches_general_mosaic( @pytest.mark.parametrize(('detsec', 'datasec', 'expected_slice'), [ ('[6:1,1:4]', '[1:6,1:4]', (slice(None), slice(None, None, -1))), ('[1:6,1:4]', '[6:1,1:4]', (slice(None), slice(None, None, -1))), + ('[1:6,4:1]', '[1:6,1:4]', (slice(None, None, -1), slice(None))), + ('[1:6,1:4]', '[1:6,4:1]', (slice(None, None, -1), slice(None))), ('[1:4,1:4]', '[2:5,1:4]', (slice(None), slice(1, 5))), ]) def test_single_component_still_crops_and_flips(detsec, datasec, expected_slice):