Avoid unnecessary copies in aligned single-component mosaics

banzai/mosaic.py

@@ -14,13 +14,29 @@ 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 = 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.
+            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,15 +47,33 @@ 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)
 
         image.primary_hdu = mosaiced_data
         return image
 
+    @staticmethod
+    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
+                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 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):
         x_detector_sections = []

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,112 @@ 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, 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=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]
+    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: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):
+    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])
+
+
+@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]'],