From af0ef972eaff561618440eac6c7e3affa9abd40c Mon Sep 17 00:00:00 2001 From: Jon Date: Wed, 15 Jul 2026 05:27:31 +0000 Subject: [PATCH 1/5] Optimize the aperture photometry operation --- .../analysis/source_catalog.py | 7 +- datalab/datalab_session/analysis/wcs.py | 8 +- .../data_operations/aperture_photometry.py | 11 +- .../tests/test_aperture_photometry.py | 142 ++++++++++-- .../datalab_session/tests/test_operations.py | 14 +- .../utils/aperture_light_curve.py | 214 +++++++++++++----- datalab/datalab_session/utils/centroiding.py | 4 +- .../datalab_session/utils/comparison_stars.py | 147 ++++++------ datalab/datalab_session/utils/file_utils.py | 15 +- .../utils/photometry_diagnostics.py | 71 ++++-- 10 files changed, 445 insertions(+), 188 deletions(-) diff --git a/datalab/datalab_session/analysis/source_catalog.py b/datalab/datalab_session/analysis/source_catalog.py index bc86210..90f4094 100644 --- a/datalab/datalab_session/analysis/source_catalog.py +++ b/datalab/datalab_session/analysis/source_catalog.py @@ -1,8 +1,7 @@ -import numpy as np from django.contrib.auth.models import User from datalab.datalab_session.exceptions import ClientAlertException -from datalab.datalab_session.utils.file_utils import get_hdu, scale_points +from datalab.datalab_session.utils.file_utils import get_fits_dimensions, get_hdu, scale_points from datalab.datalab_session.utils.filecache import FileCache from datalab.datalab_session.utils.flux_to_mag import flux_to_mag @@ -33,7 +32,6 @@ def source_catalog(input: dict, user: User): raise ClientAlertException(f"Download of {input['basename']} timed out") cat_hdu = get_hdu(file_path, 'CAT') - sci_hdu = get_hdu(file_path, 'SCI') DECIMALS_OF_PRECISION = 6 MAX_SOURCE_CATALOG_SIZE = min(len(cat_hdu.data["x"]), 1000) @@ -65,7 +63,8 @@ def source_catalog(input: dict, user: User): dec = None # scale the x_points and y_points from the fits pixel coords to the jpg coords - fits_height, fits_width = np.shape(sci_hdu.data) + # Shape comes from the SCI header: touching sci_hdu.data would decompress the whole image + fits_height, fits_width = get_fits_dimensions(file_path) x_points, y_points = scale_points(fits_height, fits_width, input['width'], input['height'], x_points=x_points, y_points=y_points) x, y = scale_points(fits_height, fits_width, input['width'], input['height'], x_points=x, y_points=y) diff --git a/datalab/datalab_session/analysis/wcs.py b/datalab/datalab_session/analysis/wcs.py index c116085..13e7bfa 100644 --- a/datalab/datalab_session/analysis/wcs.py +++ b/datalab/datalab_session/analysis/wcs.py @@ -4,7 +4,7 @@ from datalab.datalab_session.exceptions import ClientAlertException -from datalab.datalab_session.utils.file_utils import get_hdu +from datalab.datalab_session.utils.file_utils import get_fits_header from datalab.datalab_session.utils.filecache import FileCache @@ -18,8 +18,8 @@ def wcs(input: dict, user: User): """ try: file_path = FileCache().get_fits(input['basename'], input['source'], user) - sci_hdu = get_hdu(file_path, 'SCI') - fits_dimensions = [sci_hdu.header.get('NAXIS1'), sci_hdu.header.get('NAXIS2')] + sci_header = get_fits_header(file_path, 'SCI') + fits_dimensions = [sci_header.get('NAXIS1'), sci_header.get('NAXIS2')] except TimeoutError as e: raise ClientAlertException(f"Download of {input['basename']} FITs timed out") except TypeError as e: @@ -28,7 +28,7 @@ def wcs(input: dict, user: User): raise ClientAlertException(f"No FITs file found for this image") try: - wcs = WCS(sci_hdu.header) + wcs = WCS(sci_header) wcs_solution = wcs.wcs wcs_cd = wcs_solution.cd output = { diff --git a/datalab/datalab_session/data_operations/aperture_photometry.py b/datalab/datalab_session/data_operations/aperture_photometry.py index 42247f0..4b7d5e4 100644 --- a/datalab/datalab_session/data_operations/aperture_photometry.py +++ b/datalab/datalab_session/data_operations/aperture_photometry.py @@ -4,8 +4,8 @@ from django.contrib.auth.models import User from datalab.datalab_session.data_operations.data_operation import BaseDataOperation -from datalab.datalab_session.data_operations.input_data_handler import InputDataHandler from datalab.datalab_session.exceptions import ClientAlertException +from datalab.datalab_session.utils.filecache import FileCache from datalab.datalab_session.utils.format import Format from datalab.datalab_session.utils.aperture_light_curve import ( DEFAULT_ANNULUS_INNER_RADIUS, @@ -128,12 +128,15 @@ def operate(self, submitter: User): min_comparisons = int(self.input_data.get('min_comparisons', DEFAULT_MIN_COMPARISONS)) max_comparisons = int(self.input_data.get('max_comparisons', DEFAULT_MAX_COMPARISONS)) self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['INPUT_PROCESSING_PERCENTAGE_COMPLETION']) - input_handlers = [ - InputDataHandler(submitter, input_file['basename'], input_file.get('source')) + # Resolve inputs to local file-cache paths only. Pixel data is loaded (and released) + # frame by frame inside generate_light_curve, never held for all inputs at once. + file_cache = FileCache() + fits_paths = [ + file_cache.get_fits(input_file['basename'], input_file.get('source'), submitter) for input_file in input_files ] result = generate_light_curve( - input_handlers=input_handlers, + fits_paths=fits_paths, target_ra_deg=target_ra, target_dec_deg=target_dec, aperture_radius=aperture_radius, diff --git a/datalab/datalab_session/tests/test_aperture_photometry.py b/datalab/datalab_session/tests/test_aperture_photometry.py index 0cd6814..7df8624 100644 --- a/datalab/datalab_session/tests/test_aperture_photometry.py +++ b/datalab/datalab_session/tests/test_aperture_photometry.py @@ -33,11 +33,11 @@ TEST_DEG_PER_PIXEL = 1.0 / 3600.0 def print_nearest_fits_catalog_target_matches( - input_handlers: list[Any], + fits_paths: list[str], target_ra_deg: float, target_dec_deg: float, ) -> None: - frames = _validated_frame_contexts(input_handlers) + frames = _validated_frame_contexts(fits_paths) print("\nNearest FITS catalog rows to target:") for frame in frames: candidates = [ @@ -179,8 +179,8 @@ def setUp(self) -> None: def tearDown(self) -> None: self.temp_dir.cleanup() - def write_frames(self, frames: dict[str, dict[str, Any]]) -> list[Any]: - input_handlers: list[Any] = [] + def write_frames(self, frames: dict[str, dict[str, Any]]) -> list[str]: + fits_paths: list[str] = [] for name, frame in frames.items(): path = os.path.join(self.temp_dir.name, name) header = fits.Header() @@ -192,19 +192,8 @@ def write_frames(self, frames: dict[str, dict[str, Any]]) -> list[Any]: ] hdus.append(self._cat_hdu(frame["second_hdu"])) fits.HDUList(hdus).writeto(path, overwrite=True) - input_handlers.append(self.input_handler_for_path(path)) - return input_handlers - - def input_handler_for_path(self, path: str) -> Any: - with fits.open(path) as hdul: - hdus = {hdu.name: hdu.copy() for hdu in hdul} - hdus["PRIMARY"] = hdul[0].copy() - return SimpleNamespace( - fits_file=path, - sci_hdu=hdus["SCI"], - sci_data=hdus["SCI"].data, - get_hdu=lambda extension=None, hdus=hdus: hdus[extension or "SCI"], - ) + fits_paths.append(path) + return fits_paths def _cat_hdu(self, rows: list[dict[str, Any]]) -> fits.BinTableHDU: if not rows: @@ -326,7 +315,7 @@ def test_target_recenter_falls_back_to_wcs_position(self) -> None: # Point the target at an empty patch of sky (~16px from any source) so centroiding cannot # lock on. _measure_target must never raise or drift onto a neighbour: it measures at the # authoritative WCS position instead of the real source at (30.3, 28.8). - from datalab.datalab_session.utils.aperture_light_curve import _measure_target + from datalab.datalab_session.utils.aperture_light_curve import _load_frame_image, _measure_target from datalab.datalab_session.utils.fits_metadata import world_to_pixel frames, _ = build_frame_set() @@ -338,6 +327,7 @@ def test_target_recenter_falls_back_to_wcs_position(self) -> None: measurement = _measure_target( frame=frame, + image=_load_frame_image(frame.fits_path), target_ra_deg=empty_ra, target_dec_deg=empty_dec, aperture_radius=4.0, @@ -631,6 +621,111 @@ def test_determinism(self) -> None: [row.target_calibrated_apparent_magnitude for row in result2.light_curve_rows], ) + def test_pixel_data_streams_one_frame_at_a_time(self) -> None: + # The pipeline exists to keep memory flat in the input count: each frame's pixels must be + # loaded exactly once, at most one frame's pixels may be alive at any moment, and nothing + # in the result may retain pixel arrays after the run. + import gc + import weakref + from unittest import mock + + from datalab.datalab_session.utils import aperture_light_curve as light_curve_module + + frames, (target_ra, target_dec) = build_frame_set() + fits_paths = self.write_frames(frames) + + loaded_refs: list[weakref.ref] = [] + loaded_paths: list[str] = [] + max_concurrent_images = 0 + original_load = light_curve_module._load_frame_image + + def tracking_load(fits_path: str): + nonlocal max_concurrent_images + image = original_load(fits_path) + alive = sum(1 for ref in loaded_refs if ref() is not None) + 1 + max_concurrent_images = max(max_concurrent_images, alive) + loaded_refs.append(weakref.ref(image)) + loaded_paths.append(fits_path) + return image + + with mock.patch.object(light_curve_module, "_load_frame_image", new=tracking_load): + result = generate_light_curve(fits_paths, target_ra, target_dec, 4.0, 6.0, 9.0) + + gc.collect() + self.assertEqual(len(result.light_curve_rows), 3) + self.assertEqual(sorted(loaded_paths), sorted(fits_paths)) + self.assertEqual(max_concurrent_images, 1) + self.assertTrue(all(ref() is None for ref in loaded_refs)) + + def test_frame_preview_downsamples_large_frames(self) -> None: + from datalab.datalab_session.utils.photometry_diagnostics import build_frame_preview + + image = np.full((2400, 3200), 100.0, dtype=np.float32) + + preview = build_frame_preview(image, max_dimension=1500) + + self.assertEqual(preview.gray.dtype, np.uint8) + self.assertEqual(preview.gray.shape, (800, 1066)) + self.assertLessEqual(max(preview.gray.shape), 1500) + self.assertAlmostEqual(preview.scale, 1.0 / 3.0) + + def test_frame_preview_keeps_small_frames_full_resolution(self) -> None: + from datalab.datalab_session.utils.photometry_diagnostics import build_frame_preview + + preview = build_frame_preview(np.full((80, 100), 100.0, dtype=np.float32)) + + self.assertEqual(preview.gray.shape, (80, 100)) + self.assertEqual(preview.scale, 1.0) + + def test_overlay_positions_scale_to_downsampled_preview(self) -> None: + from datalab.datalab_session.utils.photometry_diagnostics import ( + build_frame_preview, + candidate_overlay_jpeg_base64, + ) + + height, width = 2400, 3200 + header = { + "CTYPE1": "RA---TAN", + "CTYPE2": "DEC--TAN", + "CUNIT1": "deg", + "CUNIT2": "deg", + "CRVAL1": 100.0, + "CRVAL2": 20.0, + "CRPIX1": 0.0, + "CRPIX2": 0.0, + "CD1_1": TEST_DEG_PER_PIXEL, + "CD1_2": 0.0, + "CD2_1": 0.0, + "CD2_2": TEST_DEG_PER_PIXEL, + } + frame = SimpleNamespace(fits_path="big.fits", header=header, width=width, height=height) + preview = build_frame_preview(np.full((height, width), 100.0, dtype=np.float32)) + self.assertLess(preview.scale, 1.0) + target_full_res = SimpleNamespace(x=600.0, y=300.0) + + encoded = candidate_overlay_jpeg_base64( + frame=frame, + preview=preview, + stars=[], + measurements=[], + target_measurement=target_full_res, + aperture_radius=4.0, + ) + + rgb = np.asarray(Image.open(BytesIO(base64.b64decode(encoded))).convert("RGB")) + self.assertEqual(rgb.shape[:2], preview.gray.shape) + orange = np.argwhere( + (rgb[:, :, 0] > 180) & + (rgb[:, :, 1] > 80) & + (rgb[:, :, 1] < 170) & + (rgb[:, :, 2] < 90) + ) + self.assertGreater(len(orange), 0) + expected_x = 600.0 * preview.scale + expected_y = (preview.gray.shape[0] - 1) - 300.0 * preview.scale + self.assertAlmostEqual(float(np.mean(orange[:, 1])), expected_x, delta=3.0) + self.assertAlmostEqual(float(np.mean(orange[:, 0])), expected_y, delta=3.0) + def test_default_fits_dependencies_read_sci_header_and_cat_rows(self) -> None: frames, (target_ra, target_dec) = build_frame_set(frame_count=1) handle = tempfile.NamedTemporaryFile(suffix=".fits", delete=False) @@ -659,9 +754,8 @@ def test_default_fits_dependencies_read_sci_header_and_cat_rows(self) -> None: ]) hdul.writeto(path, overwrite=True) try: - input_handler = self.input_handler_for_path(path) result = generate_light_curve( - [input_handler], + [path], target_ra_deg=target_ra, target_dec_deg=target_dec, aperture_radius=4.0, @@ -699,18 +793,17 @@ def test_second_hdu_requires_finite_flux_for_candidate_comparisons(self) -> None def test_real_compressed_fits_aperture_photometry_prints_diagnostics_and_results(self) -> None: fits_paths = sorted(str(path) for path in APERTURE_PHOTOMETRY_TEST_DIR.glob("*.fits.fz")) self.assertEqual(len(fits_paths), 3) - input_handlers = [self.input_handler_for_path(path) for path in fits_paths] target_ra_deg = 199.150264 target_dec_deg = 42.093592 print_nearest_fits_catalog_target_matches( - input_handlers, + fits_paths, target_ra_deg=target_ra_deg, target_dec_deg=target_dec_deg, ) result = generate_light_curve( - input_handlers, + fits_paths, target_ra_deg=target_ra_deg, target_dec_deg=target_dec_deg, aperture_radius=7.64, @@ -751,9 +844,8 @@ def test_default_fits_dependencies_reject_missing_cat(self) -> None: fits.ImageHDU(data=np.zeros((20, 20), dtype=float), header=header, name="SCI"), ]).writeto(path, overwrite=True) try: - input_handler = self.input_handler_for_path(path) with self.assertRaisesRegex(LightCurveError, "requires at least 1 valid input file"): - generate_light_curve([input_handler], 100.0, 20.0, 2.0, 3.0, 5.0) + generate_light_curve([path], 100.0, 20.0, 2.0, 3.0, 5.0) finally: os.remove(path) diff --git a/datalab/datalab_session/tests/test_operations.py b/datalab/datalab_session/tests/test_operations.py index bc1f338..0cd63fd 100644 --- a/datalab/datalab_session/tests/test_operations.py +++ b/datalab/datalab_session/tests/test_operations.py @@ -289,7 +289,7 @@ def valid_input_data(self): } @mock.patch('datalab.datalab_session.data_operations.aperture_photometry.generate_light_curve') - @mock.patch('datalab.datalab_session.data_operations.aperture_photometry.InputDataHandler') + @mock.patch('datalab.datalab_session.data_operations.aperture_photometry.FileCache') @mock.patch.object(AperturePhotometry, 'set_status') @mock.patch.object(AperturePhotometry, 'set_output') @mock.patch.object(AperturePhotometry, 'set_operation_progress') @@ -298,11 +298,10 @@ def test_operate_requires_and_passes_explicit_radii_and_filter( mock_set_operation_progress, mock_set_output, mock_set_status, - mock_input_data_handler, + mock_file_cache, mock_generate_light_curve, ): - input_handler = SimpleNamespace(fits_file='/tmp/fits_1.fits') - mock_input_data_handler.return_value = input_handler + mock_file_cache.return_value.get_fits.return_value = '/tmp/fits_1.fits' mock_generate_light_curve.return_value = SimpleNamespace( light_curve_rows=[ LightCurveRow( @@ -332,8 +331,9 @@ def test_operate_requires_and_passes_explicit_radii_and_filter( aperture_photometry = AperturePhotometry(input_data) aperture_photometry.operate(None) + mock_file_cache.return_value.get_fits.assert_called_once_with('fits_1', 'local', None) mock_generate_light_curve.assert_called_once_with( - input_handlers=[input_handler], + fits_paths=['/tmp/fits_1.fits'], target_ra_deg=10.0, target_dec_deg=20.0, aperture_radius=7.64, @@ -383,11 +383,11 @@ def test_operate_allows_missing_filter(self): del input_data['input_files'][0]['filter'] with mock.patch('datalab.datalab_session.data_operations.aperture_photometry.generate_light_curve') as mock_generate_light_curve, \ - mock.patch('datalab.datalab_session.data_operations.aperture_photometry.InputDataHandler') as mock_input_data_handler, \ + mock.patch('datalab.datalab_session.data_operations.aperture_photometry.FileCache') as mock_file_cache, \ mock.patch.object(AperturePhotometry, 'set_output') as mock_set_output, \ mock.patch.object(AperturePhotometry, 'set_operation_progress'), \ mock.patch.object(AperturePhotometry, 'set_status'): - mock_input_data_handler.return_value = SimpleNamespace(fits_file='/tmp/fits_1.fits') + mock_file_cache.return_value.get_fits.return_value = '/tmp/fits_1.fits' mock_generate_light_curve.return_value = SimpleNamespace( light_curve_rows=[], selected_comparison_stars=[], diff --git a/datalab/datalab_session/utils/aperture_light_curve.py b/datalab/datalab_session/utils/aperture_light_curve.py index 845344b..d46990d 100644 --- a/datalab/datalab_session/utils/aperture_light_curve.py +++ b/datalab/datalab_session/utils/aperture_light_curve.py @@ -6,12 +6,15 @@ from typing import Any, Iterable, Mapping, Sequence import numpy as np +from astropy.io import fits from astropy.wcs import WCS from dateutil.parser import ParserError, parse as parse_date from datalab.datalab_session.utils.comparison_stars import ( ComparisonMeasurement, ComparisonStar, + candidate_stars_from_catalog, + measure_candidate_on_frame, select_comparison_stars, ) from datalab.datalab_session.utils.centroiding import calculate_background_model, centroid @@ -29,6 +32,8 @@ minimum_angular_neighbor_distance_arcsec, ) from datalab.datalab_session.utils.photometry_diagnostics import ( + FramePreview, + build_frame_preview, candidate_overlay_jpeg_base64, comparison_star_validation_diagnostics, ) @@ -41,6 +46,16 @@ SOURCE_CATALOG_DEC_KEY = "dec" SOURCE_CATALOG_MAG_KEY = "mag" SOURCE_CATALOG_FLUX_KEY = "flux" +# The only CAT columns the pipeline reads. CAT tables carry many more columns, and whole rows kept +# per frame for the full run are a measurable share of operation memory on dense fields. +SOURCE_CATALOG_COLUMNS = ( + "id", + "name", + SOURCE_CATALOG_RA_KEY, + SOURCE_CATALOG_DEC_KEY, + SOURCE_CATALOG_MAG_KEY, + SOURCE_CATALOG_FLUX_KEY, +) EDGE_MARGIN_PX = 2.0 TARGET_PROXIMITY_FACTOR = 2.0 # A target recenter is accepted only if the centroid moves less than this many pixels from the @@ -68,12 +83,15 @@ class LightCurveError(ValueError): @dataclass(frozen=True) class FrameContext: """ - Validates FITS frame data needed by the aperture photometry pipeline. + Validated FITS frame metadata needed by the aperture photometry pipeline. + + Deliberately holds no pixel data: full-resolution images are streamed through the pixel + pass one frame at a time (see _measure_frame_pixels), so peak memory stays flat no matter + how many frames are submitted. """ fits_path: str date_obs: datetime header: Mapping[str, Any] - image: np.ndarray second_hdu_rows: tuple[Mapping[str, Any], ...] width: int height: int @@ -138,7 +156,7 @@ class LightCurveResult: def generate_light_curve( - input_handlers: list[Any], + fits_paths: list[str], target_ra_deg: float, target_dec_deg: float, aperture_radius: float, @@ -148,21 +166,25 @@ def generate_light_curve( max_comparisons: int = 10, ) -> LightCurveResult: """ - Generates a calibrated target light curve from input FITS files, using comparison stars from the source catalog. - - Validates frames, measures the target, builds a comp star catalog, - selects a comparison ensemble, and produces calibrated light curve rows with diagnostics for the frontend. + Generates a calibrated target light curve from local input FITS files, using comparison + stars from the source catalog. + + Validates frame metadata and builds the comparison-star candidate catalog from headers and + CAT tables alone, then streams pixel data one frame at a time to measure the target and + every candidate, selects a comparison ensemble, and produces calibrated light curve rows + with diagnostics for the frontend. At most one frame's full-resolution pixels are in + memory at any point, so memory does not grow with the number of input frames. """ log.info( "Aperture Photometry pipeline starting: " - f"fits_count={len(input_handlers)}, target_ra={target_ra_deg:.8f}, target_dec={target_dec_deg:.8f}, " + f"fits_count={len(fits_paths)}, target_ra={target_ra_deg:.8f}, target_dec={target_dec_deg:.8f}, " f"aperture_radius={aperture_radius:.3f}, " f"annulus_inner_radius={annulus_inner_radius:.3f}, " f"annulus_outer_radius={annulus_outer_radius:.3f}, " f"min_comparisons={min_comparisons}, max_comparisons={max_comparisons}" ) _validate_inputs( - input_handlers=input_handlers, + fits_paths=fits_paths, aperture_radius=aperture_radius, annulus_inner_radius=annulus_inner_radius, annulus_outer_radius=annulus_outer_radius, @@ -171,25 +193,50 @@ def generate_light_curve( ) diagnostics: list[str] = [] - frames = _validated_frame_contexts(input_handlers) + frames = _validated_frame_contexts(fits_paths) diagnostics_by_fits_basename: dict[str, list[str]] = { os.path.basename(frame.fits_path): [] for frame in frames } - target_measurements = { - frame.fits_path: _measure_target( + + catalog = _build_field_star_catalog( + frames=frames, + target_ra_deg=target_ra_deg, + target_dec_deg=target_dec_deg, + aperture_radius=aperture_radius, + annulus_outer_radius=annulus_outer_radius, + ) + log.info( + "Aperture Photometry comparison catalog built: " + f"valid_candidates={len(catalog)}" + ) + candidate_stars = candidate_stars_from_catalog(catalog) + + target_measurements: dict[str, TargetMeasurement] = {} + previews: dict[str, FramePreview] = {} + measurements_by_candidate: dict[str, dict[str, ComparisonMeasurement]] = { + candidate.candidate_id: {} for candidate in candidate_stars + } + failed_candidate_ids: set[str] = set() + for frame in frames: + target, frame_measurements, newly_failed, preview = _measure_frame_pixels( frame=frame, + candidate_stars=candidate_stars, + skip_candidate_ids=failed_candidate_ids, target_ra_deg=target_ra_deg, target_dec_deg=target_dec_deg, aperture_radius=aperture_radius, annulus_inner_radius=annulus_inner_radius, annulus_outer_radius=annulus_outer_radius, ) - for frame in frames - } - for frame in frames: - target = target_measurements[frame.fits_path] + target_measurements[frame.fits_path] = target + previews[frame.fits_path] = preview + failed_candidate_ids |= newly_failed + for candidate_id in newly_failed: + measurements_by_candidate.pop(candidate_id, None) + for candidate_id, measurement in frame_measurements.items(): + measurements_by_candidate[candidate_id][frame.fits_path] = measurement log.info( "Aperture Photometry target measurement: " f"frame={frame.fits_path}, centroid=({target.x:.3f}, {target.y:.3f}), " @@ -197,26 +244,13 @@ def generate_light_curve( f"background={target.mean_background_per_pixel:.6f}, peak={target.peak_pixel_value:.6f}" ) - catalog = _build_field_star_catalog( - frames=frames, - target_ra_deg=target_ra_deg, - target_dec_deg=target_dec_deg, - aperture_radius=aperture_radius, - annulus_outer_radius=annulus_outer_radius, - ) - log.info( - "Aperture Photometry comparison catalog built: " - f"valid_candidates={len(catalog)}" - ) target_mag_proxy = _target_magnitude_proxy(target_measurements.values()) log.info(f"Aperture Photometry target magnitude proxy: {target_mag_proxy:.6f}") selection = select_comparison_stars( frames=frames, - catalog=catalog, + candidates=candidate_stars, + measurements_by_candidate=measurements_by_candidate, target_mag_proxy=target_mag_proxy, - aperture_radius=aperture_radius, - annulus_inner_radius=annulus_inner_radius, - annulus_outer_radius=annulus_outer_radius, min_comparisons=min_comparisons, max_comparisons=max_comparisons, error_class=LightCurveError, @@ -313,6 +347,7 @@ def generate_light_curve( diagnostics_by_fits_basename[os.path.basename(frame.fits_path)].extend(frame_diagnostics) diagnostic_images_by_fits_basename[os.path.basename(frame.fits_path)] = candidate_overlay_jpeg_base64( frame=frame, + preview=previews[frame.fits_path], stars=selection.selected_stars, measurements=comparison_measurements, target_measurement=target, @@ -361,15 +396,15 @@ def generate_light_curve( def _validate_inputs( *, - input_handlers: Sequence[Any], + fits_paths: Sequence[str], aperture_radius: float, annulus_inner_radius: float, annulus_outer_radius: float, min_comparisons: int, max_comparisons: int, ) -> None: - if not input_handlers: - raise LightCurveError("input_handlers must be a non-empty list.") + if not fits_paths: + raise LightCurveError("fits_paths must be a non-empty list.") if aperture_radius <= 0: raise LightCurveError("aperture_radius must be > 0.") if annulus_inner_radius <= aperture_radius: @@ -380,17 +415,27 @@ def _validate_inputs( raise LightCurveError("min_comparisons and max_comparisons must be positive and min_comparisons <= max_comparisons.") -def _validated_frame_contexts(input_handlers: Sequence[Any]) -> list[FrameContext]: +def _validated_frame_contexts(fits_paths: Sequence[str]) -> list[FrameContext]: + """ + Builds validated frame metadata for each input FITS path. + + Reads only the SCI header and the CAT table -- never SCI pixel data -- so validation memory + and time stay flat regardless of frame count or sensor size. Frames that fail validation + are ignored with a warning. + """ frames: list[FrameContext] = [] - for input_handler in input_handlers: - fits_path = input_handler.fits_file + for fits_path in fits_paths: log.info(f"Aperture Photometry validating FITS frame: {fits_path}") try: - image = np.asarray(input_handler.sci_data, dtype=float) - if image.ndim != 2: + with fits.open(fits_path) as hdul: + header = dict(hdul["SCI"].header) + second_hdu_rows = tuple(_cat_rows(hdul["CAT"].data)) + + if int(header.get("NAXIS", 0)) != 2: raise LightCurveError(f"Primary image for {fits_path} is not a 2D array.") + width = int(header["NAXIS1"]) + height = int(header["NAXIS2"]) - header = dict(input_handler.sci_hdu.header) date_obs_value = header.get("DATE-OBS") if not isinstance(date_obs_value, str) or not date_obs_value.strip(): raise LightCurveError(f"Missing DATE-OBS in {fits_path}.") @@ -400,26 +445,24 @@ def _validated_frame_contexts(input_handlers: Sequence[Any]) -> list[FrameContex raise LightCurveError(f"Malformed DATE-OBS in {fits_path}: {date_obs_value!r}") from exc if date_obs.tzinfo is None: date_obs = date_obs.replace(tzinfo=timezone.utc) - second_hdu_rows = tuple(_cat_rows(input_handler.get_hdu("CAT").data)) if not second_hdu_rows: raise LightCurveError(f"Second HDU is missing or empty for {fits_path}.") - _validate_wcs(header, fits_path, image.shape) + _validate_wcs(header, fits_path, (height, width)) _validate_second_hdu(second_hdu_rows, fits_path) log.info( "Aperture Photometry frame validated: " f"frame={fits_path}, date_obs={date_obs.isoformat()}, " - f"image_shape={image.shape}, catalog_rows={len(second_hdu_rows)}" + f"image_shape={(height, width)}, catalog_rows={len(second_hdu_rows)}" ) frames.append( FrameContext( fits_path=fits_path, date_obs=date_obs, header=header, - image=image, second_hdu_rows=second_hdu_rows, - width=int(image.shape[1]), - height=int(image.shape[0]), + width=width, + height=height, ) ) except LightCurveError as exc: @@ -444,10 +487,76 @@ def _validated_frame_contexts(input_handlers: Sequence[Any]) -> list[FrameContex return frames +def _load_frame_image(fits_path: str) -> np.ndarray: + """ + Loads one frame's SCI pixel data as float32. + + float32 matches the archive's native SCI pixel type; asking for float64 here would double + every frame's in-memory size (photometry sums already accumulate in double precision). + """ + with fits.open(fits_path, memmap=False) as hdul: + image = np.asarray(hdul["SCI"].data, dtype=np.float32) + if image.ndim != 2: + raise LightCurveError(f"Primary image for {fits_path} is not a 2D array.") + return image + + +def _measure_frame_pixels( + *, + frame: FrameContext, + candidate_stars: Sequence[ComparisonStar], + skip_candidate_ids: set[str], + target_ra_deg: float, + target_dec_deg: float, + aperture_radius: float, + annulus_inner_radius: float, + annulus_outer_radius: float, +) -> tuple[TargetMeasurement, dict[str, ComparisonMeasurement], set[str], FramePreview]: + """ + Runs all pixel-dependent work for one frame: the target measurement, a measurement of every + comparison candidate (minus skip_candidate_ids), and the downsampled diagnostic preview. + + The full-resolution image exists only inside this function, so it is released before the + caller moves on to the next frame. + + Returns the target measurement, this frame's candidate measurements by candidate_id, the + ids of candidates that failed to measure on this frame, and the preview. + """ + image = _load_frame_image(frame.fits_path) + target_measurement = _measure_target( + frame=frame, + image=image, + target_ra_deg=target_ra_deg, + target_dec_deg=target_dec_deg, + aperture_radius=aperture_radius, + annulus_inner_radius=annulus_inner_radius, + annulus_outer_radius=annulus_outer_radius, + ) + candidate_measurements: dict[str, ComparisonMeasurement] = {} + failed_candidate_ids: set[str] = set() + for candidate in candidate_stars: + if candidate.candidate_id in skip_candidate_ids: + continue + try: + candidate_measurements[candidate.candidate_id] = measure_candidate_on_frame( + frame=frame, + image=image, + candidate=candidate, + aperture_radius=aperture_radius, + annulus_inner_radius=annulus_inner_radius, + annulus_outer_radius=annulus_outer_radius, + error_class=LightCurveError, + ) + except LightCurveError: + failed_candidate_ids.add(candidate.candidate_id) + preview = build_frame_preview(image) + return target_measurement, candidate_measurements, failed_candidate_ids, preview + + def _cat_rows(data: Any) -> list[dict[str, Any]]: if data is None: return [] - names = list(data.names or []) + names = [name for name in (data.names or []) if name in SOURCE_CATALOG_COLUMNS] return [ { name: data[name][index].item() if hasattr(data[name][index], "item") else data[name][index] @@ -488,6 +597,7 @@ def _validate_second_hdu(rows: Sequence[Mapping[str, Any]], fits_path: str) -> N def _measure_target( *, frame: FrameContext, + image: np.ndarray, target_ra_deg: float, target_dec_deg: float, aperture_radius: float, @@ -495,7 +605,9 @@ def _measure_target( annulus_outer_radius: float, ) -> TargetMeasurement: """ - Converts the target RA and Dec to pixel coordinates, centroids the source, and measures aperture photometry. + Converts the target RA and Dec to pixel coordinates, centroids the source, and measures + aperture photometry. image is the frame's pixel data, passed separately from the metadata + so the streaming pixel pass controls how long it stays in memory. The target is never allowed to drop a frame: if centroiding fails or the refinement drifts too far from the WCS position, it measures at the authoritative WCS position instead. @@ -516,7 +628,7 @@ def _measure_target( ) centroid_result = centroid( - image=frame.image, + image=image, x_click=initial_x, y_click=initial_y, radius=aperture_radius_px, @@ -535,7 +647,7 @@ def _measure_target( # Re-estimate the background at the WCS position: a drifted annulus can straddle the host # galaxy and bias the sky level, which is exactly the pull we are rejecting. background_model = calculate_background_model( - frame.image, + image, x_center, y_center, aperture_radius_px, @@ -560,7 +672,7 @@ def _measure_target( ) photometry = measure_aperture( - image=frame.image, + image=image, x_center=x_center, y_center=y_center, aperture_radius_px=aperture_radius_px, diff --git a/datalab/datalab_session/utils/centroiding.py b/datalab/datalab_session/utils/centroiding.py index 6036814..dcf6944 100644 --- a/datalab/datalab_session/utils/centroiding.py +++ b/datalab/datalab_session/utils/centroiding.py @@ -257,7 +257,9 @@ def centroid( Returns a CentroidResult containing the centroid position, localbackground estimate, peak value, and any relevant messages. """ - image = np.asarray(image, dtype=float) + # No dtype here: forcing float64 would copy the entire frame on every call, and this runs once + # per comparison candidate per frame. Pixels are read as Python floats, so any numeric dtype works. + image = np.asarray(image) x_center = x_click y_center = y_click radius = max(radius, 3.0) diff --git a/datalab/datalab_session/utils/comparison_stars.py b/datalab/datalab_session/utils/comparison_stars.py index 725a7b4..96cd44d 100644 --- a/datalab/datalab_session/utils/comparison_stars.py +++ b/datalab/datalab_session/utils/comparison_stars.py @@ -61,29 +61,26 @@ class ComparisonSelectionResult: def select_comparison_stars( *, frames: Sequence[Any], - catalog: Sequence[dict[str, Any]], + candidates: Sequence[ComparisonStar], + measurements_by_candidate: Mapping[str, Mapping[str, ComparisonMeasurement]], target_mag_proxy: float, - aperture_radius: float, - annulus_inner_radius: float, - annulus_outer_radius: float, min_comparisons: int, max_comparisons: int, error_class: type[Exception] = ValueError, ) -> ComparisonSelectionResult: """ - Selects comparison stars from source catalog candidates. + Selects comparison stars from source catalog candidates using their per-frame measurements. - Measures every candidate, drops variable stars and zero-point-inconsistent (blended or - mismatched) catalog matches, then ranks the rest by how close their measured brightness is - to the target's. Returns the comparison ensemble for calibration. + Measurement happens upstream (one frame's pixels at a time); this is pure math on the + collected measurements. Drops candidates without a valid positive measurement on every + frame, then variable stars and zero-point-inconsistent (blended or mismatched) catalog + matches, and ranks the rest by how close their measured brightness is to the target's. + Returns the comparison ensemble for calibration. """ - enriched, measurements_by_candidate = _measure_and_rank_candidates( + enriched = _rank_measured_candidates( frames=frames, - catalog=catalog, - aperture_radius=aperture_radius, - annulus_inner_radius=annulus_inner_radius, - annulus_outer_radius=annulus_outer_radius, - error_class=error_class, + candidates=candidates, + measurements_by_candidate=measurements_by_candidate, ) stable = [candidate for candidate in enriched if candidate.variability_score <= MAX_ACCEPTABLE_VARIABILITY] consistent = _reject_zero_point_outliers(stable) @@ -100,7 +97,7 @@ def select_comparison_stars( selected_stars=selected_stars, diagnostics=tuple(), measurements_by_candidate={ - star.candidate_id: measurements_by_candidate[star.candidate_id] + star.candidate_id: dict(measurements_by_candidate[star.candidate_id]) for star in selected_stars }, ) @@ -145,6 +142,7 @@ def _source_catalog_sort_key(candidate: ComparisonStar, target_mag_proxy: float) def measure_candidate_on_frame( *, frame: Any, + image: np.ndarray, candidate: ComparisonStar, aperture_radius: float, annulus_inner_radius: float, @@ -154,6 +152,8 @@ def measure_candidate_on_frame( """ Measures aperture photometry for one comparison-star candidate on a single FITS frame. + image is the frame's full-resolution pixel data, passed separately from the frame metadata. + Converts the candidate's RA/Dec to pixel coordinates via the frame WCS, centroids around that position to refine it (correcting small WCS or catalog inaccuracies), then measures aperture photometry at the refined position, estimating the background, summing the @@ -166,7 +166,7 @@ def measure_candidate_on_frame( annulus_outer_radius_px = arcsec_to_pixels(frame.header, annulus_outer_radius) x, y = world_to_pixel(frame.header, candidate.ra_deg, candidate.dec_deg) centroid_result = centroid( - image=frame.image, + image=image, x_click=x, y_click=y, radius=aperture_radius_px, @@ -176,7 +176,7 @@ def measure_candidate_on_frame( if not centroid_result.success: raise error_class(f"Selected comparison-star centroiding failed for {frame.fits_path}, {candidate.candidate_id}.") photometry = measure_aperture( - image=frame.image, + image=image, x_center=centroid_result.x, y_center=centroid_result.y, aperture_radius_px=aperture_radius_px, @@ -200,89 +200,82 @@ def measure_candidate_on_frame( ) -def _measure_and_rank_candidates( +def candidate_stars_from_catalog(catalog: Sequence[dict[str, Any]]) -> list[ComparisonStar]: + """ + Builds not-yet-measured ComparisonStar candidates from field-star catalog rows, ordered by + candidate_id so measurement and ranking stay deterministic. + """ + candidates: list[ComparisonStar] = [] + for candidate in sorted(catalog, key=lambda row: row["candidate_id"]): + candidates.append( + ComparisonStar( + candidate_id=candidate["candidate_id"], + ra_deg=candidate["ra_deg"], + dec_deg=candidate["dec_deg"], + reference_magnitude=float(candidate.get("reference_magnitude", candidate["second_hdu_magnitude"])), + reference_magnitude_source=str(candidate.get("reference_magnitude_source", "second_hdu")), + source_catalog_by_frame=candidate.get("source_catalog_by_frame", {}), + variability_score=math.inf, + isolation_arcsec=candidate["isolation_arcsec"], + target_separation_px=candidate["target_separation_px"], + ) + ) + return candidates + + +def _rank_measured_candidates( *, frames: Sequence[Any], - catalog: Sequence[dict[str, Any]], - aperture_radius: float, - annulus_inner_radius: float, - annulus_outer_radius: float, - error_class: type[Exception], -) -> tuple[list[ComparisonStar], dict[str, dict[str, ComparisonMeasurement]]]: + candidates: Sequence[ComparisonStar], + measurements_by_candidate: Mapping[str, Mapping[str, ComparisonMeasurement]], +) -> list[ComparisonStar]: """ - Measures each candidate across all frames and calculates variability scores. + Scores measured candidates from their per-frame measurements. - Returns the comp star candidates that have valid positive measurements across all frames, - together with those per-frame measurements keyed candidate_id -> fits_path -> measurement - so the caller can reuse them instead of re-measuring. + Keeps only candidates with a valid positive measurement on every frame and fills in their + variability scores and median measured instrumental magnitudes. """ - measured_candidates: list[tuple[dict[str, Any], ComparisonStar, np.ndarray, list[ComparisonMeasurement]]] = [] - for candidate in sorted(catalog, key=lambda row: row["candidate_id"]): - reference_magnitude = float(candidate.get("reference_magnitude", candidate["second_hdu_magnitude"])) - reference_magnitude_source = str(candidate.get("reference_magnitude_source", "second_hdu")) - source_catalog_by_frame = candidate.get("source_catalog_by_frame", {}) - candidate_star = ComparisonStar( - candidate_id=candidate["candidate_id"], - ra_deg=candidate["ra_deg"], - dec_deg=candidate["dec_deg"], - reference_magnitude=reference_magnitude, - reference_magnitude_source=reference_magnitude_source, - source_catalog_by_frame=source_catalog_by_frame, - variability_score=math.inf, - isolation_arcsec=candidate["isolation_arcsec"], - target_separation_px=candidate["target_separation_px"], - ) - try: - per_frame = [ - measure_candidate_on_frame( - frame=frame, - candidate=candidate_star, - aperture_radius=aperture_radius, - annulus_inner_radius=annulus_inner_radius, - annulus_outer_radius=annulus_outer_radius, - error_class=error_class, - ) - for frame in frames - ] - except error_class: + measured_candidates: list[tuple[ComparisonStar, np.ndarray]] = [] + for candidate in candidates: + per_frame = measurements_by_candidate.get(candidate.candidate_id, {}) + if any(frame.fits_path not in per_frame for frame in frames): continue - counts = np.asarray([measurement.net_source_counts for measurement in per_frame], dtype=float) + counts = np.asarray( + [per_frame[frame.fits_path].net_source_counts for frame in frames], + dtype=float, + ) if np.any(~np.isfinite(counts)) or np.any(counts <= 0.0): continue instrumental_mags = -2.5 * np.log10(counts) - measured_candidates.append((candidate, candidate_star, instrumental_mags, per_frame)) + measured_candidates.append((candidate, instrumental_mags)) if not measured_candidates: - return [], {} + return [] - instrumental_mag_matrix = np.vstack([row[2] for row in measured_candidates]) + instrumental_mag_matrix = np.vstack([mags for _, mags in measured_candidates]) if len(measured_candidates) > 1: frame_offsets = np.median(instrumental_mag_matrix, axis=0) variability_mag_matrix = instrumental_mag_matrix - frame_offsets else: variability_mag_matrix = instrumental_mag_matrix - selected: list[ComparisonStar] = [] - measurements_by_candidate: dict[str, dict[str, ComparisonMeasurement]] = {} - for (candidate, candidate_star, instrumental_mags, per_frame), variability_mags in zip( + scored: list[ComparisonStar] = [] + for (candidate, instrumental_mags), variability_mags in zip( measured_candidates, variability_mag_matrix, ): - selected.append( + scored.append( ComparisonStar( - candidate_id=candidate["candidate_id"], - ra_deg=candidate["ra_deg"], - dec_deg=candidate["dec_deg"], - reference_magnitude=candidate_star.reference_magnitude, - reference_magnitude_source=candidate_star.reference_magnitude_source, - source_catalog_by_frame=candidate_star.source_catalog_by_frame, + candidate_id=candidate.candidate_id, + ra_deg=candidate.ra_deg, + dec_deg=candidate.dec_deg, + reference_magnitude=candidate.reference_magnitude, + reference_magnitude_source=candidate.reference_magnitude_source, + source_catalog_by_frame=candidate.source_catalog_by_frame, variability_score=float(np.std(variability_mags)), - isolation_arcsec=candidate_star.isolation_arcsec, - target_separation_px=candidate_star.target_separation_px, + isolation_arcsec=candidate.isolation_arcsec, + target_separation_px=candidate.target_separation_px, measured_instrumental_magnitude=float(np.median(instrumental_mags)), ) ) - measurements_by_candidate[candidate["candidate_id"]] = { - measurement.fits_path: measurement for measurement in per_frame - } - return selected, measurements_by_candidate + return scored diff --git a/datalab/datalab_session/utils/file_utils.py b/datalab/datalab_session/utils/file_utils.py index 5f10f8c..03adfb3 100644 --- a/datalab/datalab_session/utils/file_utils.py +++ b/datalab/datalab_session/utils/file_utils.py @@ -20,15 +20,28 @@ def get_hdu(path: str, extension: str = 'SCI', use_fsspec: bool = False) -> list """ Returns a HDU for the fits in the given path Warning: this function returns an opened file that must be closed after use + Warning: the HDU copy eagerly loads the extension's data - decompressing the full image for + compressed extensions. If you only need the header, use get_fits_header instead. """ with fits.open(path, use_fsspec=use_fsspec) as hdu: try: extension_copy = hdu[extension].copy() except KeyError: raise ClientAlertException(f"{extension} Header not found in fits file at {path.split('/')[-1]}") - + return extension_copy +def get_fits_header(path: str, extension: str = 'SCI') -> fits.Header: + """ + Returns the header for an extension without touching its data, so large (compressed) images + are never decompressed into memory. + """ + with fits.open(path) as hdu: + try: + return hdu[extension].header.copy() + except KeyError: + raise ClientAlertException(f"{extension} Header not found in fits file at {path.split('/')[-1]}") + def get_fits_dimensions(fits_file, extension: str = 'SCI') -> tuple: with fits.open(fits_file) as hdu: hdu_shape = hdu[extension].shape diff --git a/datalab/datalab_session/utils/photometry_diagnostics.py b/datalab/datalab_session/utils/photometry_diagnostics.py index bbb3007..ff8288e 100644 --- a/datalab/datalab_session/utils/photometry_diagnostics.py +++ b/datalab/datalab_session/utils/photometry_diagnostics.py @@ -2,6 +2,7 @@ import base64 import math +from dataclasses import dataclass from io import BytesIO from typing import Any, Sequence @@ -13,22 +14,66 @@ COMPARISON_STAR_COLOR = (0, 173, 239) TARGET_COLOR = (243, 131, 33) +# Diagnostic overlays are rendered on a downsampled preview no larger than this on either side. +PREVIEW_MAX_DIMENSION = 2000 + + +@dataclass(frozen=True) +class FramePreview: + """ + Downsampled display rendering of a frame, captured while the frame's pixels were loaded. + + gray is uint8, display-oriented (y-flipped), block-mean downsampled so that preview + coordinates = full-resolution coordinates * scale. Diagnostic overlays are drawn on this + preview, so full-resolution pixels never have to be reloaded or retained for rendering. + """ + gray: np.ndarray + scale: float + + @property + def height(self) -> int: + return int(self.gray.shape[0]) + + @property + def width(self) -> int: + return int(self.gray.shape[1]) + + +def build_frame_preview(image: np.ndarray, max_dimension: int = PREVIEW_MAX_DIMENSION) -> FramePreview: + """ + Builds the downsampled uint8 display preview for a frame. + + Blocks are mean-combined (point sampling would drop stars smaller than the sampling step), + and the display stretch is computed on the downsampled array, so the only full-resolution + temporary is the trimmed copy the block reshape may make. + """ + height, width = image.shape + step = max(1, math.ceil(max(height, width) / max_dimension)) + if step > 1: + trimmed = image[: (height // step) * step, : (width // step) * step] + blocks = trimmed.reshape(height // step, step, width // step, step) + small = blocks.mean(axis=(1, 3), dtype=np.float64) + else: + small = np.asarray(image, dtype=float) + gray = np.flip(_stretch_to_uint8(small), axis=0) + return FramePreview(gray=np.ascontiguousarray(gray), scale=1.0 / step) def candidate_overlay_jpeg_base64( *, frame: Any, + preview: FramePreview, stars: Sequence[Any], measurements: Sequence[Any], target_measurement: Any, aperture_radius: float, ) -> str: - image = _normalize_image_for_jpeg(frame.image) + image = Image.fromarray(preview.gray).convert("RGB") draw = ImageDraw.Draw(image) - font = _diagnostic_overlay_font(frame.width, frame.height) + font = _diagnostic_overlay_font(preview.width, preview.height) stars_by_id = {star.candidate_id: star for star in stars} - min_dimension = max(min(frame.width, frame.height), 1) - aperture_radius_px = arcsec_to_pixels(frame.header, aperture_radius) + min_dimension = max(min(preview.width, preview.height), 1) + aperture_radius_px = arcsec_to_pixels(frame.header, aperture_radius) * preview.scale radius = max(aperture_radius_px, min_dimension * 0.018, 14.0) line_width = max(3, int(round(min_dimension * 0.004))) label_padding = max(3, int(round(min_dimension * 0.004))) @@ -36,8 +81,8 @@ def candidate_overlay_jpeg_base64( for measurement in measurements: if measurement.candidate_id not in stars_by_id: continue - x = float(measurement.x) - y = _display_y(float(measurement.y), frame.height) + x = float(measurement.x) * preview.scale + y = _display_y(float(measurement.y) * preview.scale, preview.height) if not math.isfinite(x) or not math.isfinite(y): continue @@ -51,10 +96,10 @@ def candidate_overlay_jpeg_base64( label_bbox = draw.textbbox((label_x, label_y), label, font=font) label_width = label_bbox[2] - label_bbox[0] label_height = label_bbox[3] - label_bbox[1] - if label_x + label_width + label_padding > frame.width: + if label_x + label_width + label_padding > preview.width: label_x = max(x - radius - label_width - label_padding, 0) if label_y < 0: - label_y = min(y + radius + label_padding, max(frame.height - label_height - label_padding, 0)) + label_y = min(y + radius + label_padding, max(preview.height - label_height - label_padding, 0)) draw.text( (label_x, label_y), label, @@ -62,8 +107,8 @@ def candidate_overlay_jpeg_base64( font=font, ) - target_x = float(target_measurement.x) - target_y = _display_y(float(target_measurement.y), frame.height) + target_x = float(target_measurement.x) * preview.scale + target_y = _display_y(float(target_measurement.y) * preview.scale, preview.height) if math.isfinite(target_x) and math.isfinite(target_y): target_bbox = ( target_x - radius, @@ -112,7 +157,7 @@ def comparison_star_validation_diagnostics( return diagnostics -def _normalize_image_for_jpeg(image_data: np.ndarray) -> Image.Image: +def _stretch_to_uint8(image_data: np.ndarray) -> np.ndarray: finite = np.asarray(image_data, dtype=float) finite_values = finite[np.isfinite(finite)] if finite_values.size: @@ -127,9 +172,7 @@ def _normalize_image_for_jpeg(image_data: np.ndarray) -> Image.Image: scaled = np.clip((finite - zmin) / (zmax - zmin), 0.0, 1.0) scaled = np.nan_to_num(scaled, nan=0.0, posinf=1.0, neginf=0.0) - gray = (scaled * 255.0).astype(np.uint8) - gray = np.flip(gray, axis=0) - return Image.fromarray(gray).convert("RGB") + return (scaled * 255.0).astype(np.uint8) def _display_y(y: float, height: int) -> float: From 74c84b26eb09412dd79216d4c72cec608470564b Mon Sep 17 00:00:00 2001 From: Steve Foale Date: Thu, 16 Jul 2026 06:45:42 +0100 Subject: [PATCH 2/5] Cache per-frame WCS/aperture geometry instead of rebuilding per candidate. --- .../tests/test_aperture_photometry.py | 6 +-- .../utils/aperture_light_curve.py | 29 +++++++------ .../datalab_session/utils/comparison_stars.py | 16 +++---- .../datalab_session/utils/fits_metadata.py | 43 +++++++++++++++++++ 4 files changed, 68 insertions(+), 26 deletions(-) diff --git a/datalab/datalab_session/tests/test_aperture_photometry.py b/datalab/datalab_session/tests/test_aperture_photometry.py index 7df8624..5436665 100644 --- a/datalab/datalab_session/tests/test_aperture_photometry.py +++ b/datalab/datalab_session/tests/test_aperture_photometry.py @@ -316,7 +316,7 @@ def test_target_recenter_falls_back_to_wcs_position(self) -> None: # lock on. _measure_target must never raise or drift onto a neighbour: it measures at the # authoritative WCS position instead of the real source at (30.3, 28.8). from datalab.datalab_session.utils.aperture_light_curve import _load_frame_image, _measure_target - from datalab.datalab_session.utils.fits_metadata import world_to_pixel + from datalab.datalab_session.utils.fits_metadata import frame_geometry, world_to_pixel frames, _ = build_frame_set() header = frames["frame_1.fits"]["header"] @@ -328,11 +328,9 @@ def test_target_recenter_falls_back_to_wcs_position(self) -> None: measurement = _measure_target( frame=frame, image=_load_frame_image(frame.fits_path), + geometry=frame_geometry(frame.header, 4.0, 6.0, 9.0), target_ra_deg=empty_ra, target_dec_deg=empty_dec, - aperture_radius=4.0, - annulus_inner_radius=6.0, - annulus_outer_radius=9.0, ) self.assertAlmostEqual(measurement.x, initial_x, delta=1.0e-6) diff --git a/datalab/datalab_session/utils/aperture_light_curve.py b/datalab/datalab_session/utils/aperture_light_curve.py index d46990d..f73ca5b 100644 --- a/datalab/datalab_session/utils/aperture_light_curve.py +++ b/datalab/datalab_session/utils/aperture_light_curve.py @@ -19,8 +19,10 @@ ) from datalab.datalab_session.utils.centroiding import calculate_background_model, centroid from datalab.datalab_session.utils.fits_metadata import ( + FrameGeometry, arcsec_to_pixels, frame_gain, + frame_geometry, frame_read_noise, optional_float, world_to_pixel, @@ -523,14 +525,16 @@ def _measure_frame_pixels( ids of candidates that failed to measure on this frame, and the preview. """ image = _load_frame_image(frame.fits_path) + # Build the frame's WCS and pixel-space aperture radii once, then reuse them for the target and + # every candidate. These are frame constants, so recomputing them per candidate (as the old + # arcsec_to_pixels/world_to_pixel calls did) just re-parsed the header WCS thousands of times. + geometry = frame_geometry(frame.header, aperture_radius, annulus_inner_radius, annulus_outer_radius) target_measurement = _measure_target( frame=frame, image=image, + geometry=geometry, target_ra_deg=target_ra_deg, target_dec_deg=target_dec_deg, - aperture_radius=aperture_radius, - annulus_inner_radius=annulus_inner_radius, - annulus_outer_radius=annulus_outer_radius, ) candidate_measurements: dict[str, ComparisonMeasurement] = {} failed_candidate_ids: set[str] = set() @@ -541,10 +545,8 @@ def _measure_frame_pixels( candidate_measurements[candidate.candidate_id] = measure_candidate_on_frame( frame=frame, image=image, + geometry=geometry, candidate=candidate, - aperture_radius=aperture_radius, - annulus_inner_radius=annulus_inner_radius, - annulus_outer_radius=annulus_outer_radius, error_class=LightCurveError, ) except LightCurveError: @@ -598,28 +600,27 @@ def _measure_target( *, frame: FrameContext, image: np.ndarray, + geometry: FrameGeometry, target_ra_deg: float, target_dec_deg: float, - aperture_radius: float, - annulus_inner_radius: float, - annulus_outer_radius: float, ) -> TargetMeasurement: """ Converts the target RA and Dec to pixel coordinates, centroids the source, and measures aperture photometry. image is the frame's pixel data, passed separately from the metadata - so the streaming pixel pass controls how long it stays in memory. + so the streaming pixel pass controls how long it stays in memory. geometry carries the + frame's cached WCS and pixel-space aperture radii. The target is never allowed to drop a frame: if centroiding fails or the refinement drifts too far from the WCS position, it measures at the authoritative WCS position instead. Returns the target measurement for a single frame. """ - aperture_radius_px = arcsec_to_pixels(frame.header, aperture_radius) - annulus_inner_radius_px = arcsec_to_pixels(frame.header, annulus_inner_radius) - annulus_outer_radius_px = arcsec_to_pixels(frame.header, annulus_outer_radius) + aperture_radius_px = geometry.aperture_radius_px + annulus_inner_radius_px = geometry.annulus_inner_radius_px + annulus_outer_radius_px = geometry.annulus_outer_radius_px try: - initial_x, initial_y = world_to_pixel(frame.header, target_ra_deg, target_dec_deg) + initial_x, initial_y = geometry.world_to_pixel(target_ra_deg, target_dec_deg) except Exception as exc: raise LightCurveError(f"Target WCS localization failed for {frame.fits_path}.") from exc log.info( diff --git a/datalab/datalab_session/utils/comparison_stars.py b/datalab/datalab_session/utils/comparison_stars.py index 96cd44d..70d256e 100644 --- a/datalab/datalab_session/utils/comparison_stars.py +++ b/datalab/datalab_session/utils/comparison_stars.py @@ -7,7 +7,7 @@ import numpy as np from datalab.datalab_session.utils.centroiding import centroid -from datalab.datalab_session.utils.fits_metadata import arcsec_to_pixels, frame_gain, frame_read_noise, world_to_pixel +from datalab.datalab_session.utils.fits_metadata import FrameGeometry, frame_gain, frame_read_noise from datalab.datalab_session.utils.photometry import measure_aperture @@ -143,16 +143,16 @@ def measure_candidate_on_frame( *, frame: Any, image: np.ndarray, + geometry: FrameGeometry, candidate: ComparisonStar, - aperture_radius: float, - annulus_inner_radius: float, - annulus_outer_radius: float, error_class: type[Exception] = ValueError, ) -> ComparisonMeasurement: """ Measures aperture photometry for one comparison-star candidate on a single FITS frame. image is the frame's full-resolution pixel data, passed separately from the frame metadata. + geometry carries the frame's cached WCS and pixel-space aperture radii, shared across every + candidate on the frame. Converts the candidate's RA/Dec to pixel coordinates via the frame WCS, centroids around that position to refine it (correcting small WCS or catalog inaccuracies), then measures @@ -161,10 +161,10 @@ def measure_candidate_on_frame( Returns the comparison-star measurement for this frame. """ - aperture_radius_px = arcsec_to_pixels(frame.header, aperture_radius) - annulus_inner_radius_px = arcsec_to_pixels(frame.header, annulus_inner_radius) - annulus_outer_radius_px = arcsec_to_pixels(frame.header, annulus_outer_radius) - x, y = world_to_pixel(frame.header, candidate.ra_deg, candidate.dec_deg) + aperture_radius_px = geometry.aperture_radius_px + annulus_inner_radius_px = geometry.annulus_inner_radius_px + annulus_outer_radius_px = geometry.annulus_outer_radius_px + x, y = geometry.world_to_pixel(candidate.ra_deg, candidate.dec_deg) centroid_result = centroid( image=image, x_click=x, diff --git a/datalab/datalab_session/utils/fits_metadata.py b/datalab/datalab_session/utils/fits_metadata.py index 934cea5..d1863ae 100644 --- a/datalab/datalab_session/utils/fits_metadata.py +++ b/datalab/datalab_session/utils/fits_metadata.py @@ -1,4 +1,5 @@ import math +from dataclasses import dataclass from typing import Any, Mapping import numpy as np @@ -64,3 +65,45 @@ def arcsec_to_pixels(header: Mapping[str, Any], angular_radius_arcsec: float) -> Convert an angular aperture radius to pixels for one frame, using the frame's plate scale. """ return float(angular_radius_arcsec) / pixel_scale_arcsec(header) + + +@dataclass(frozen=True) +class FrameGeometry: + """ + Per-frame WCS and pixel-space aperture geometry, built once and reused for the target and + every comparison candidate on the frame. + + Constructing a WCS from a header costs on the order of ~10 ms; arcsec_to_pixels and + world_to_pixel each built one per call, so measuring a frame's candidates rebuilt it + thousands of times. The plate scale and WCS are frame constants, so they are computed once + here instead of per candidate. + """ + wcs: WCS + aperture_radius_px: float + annulus_inner_radius_px: float + annulus_outer_radius_px: float + + def world_to_pixel(self, ra_deg: float, dec_deg: float) -> tuple[float, float]: + """Pixel coordinates of a sky position using the cached WCS (no header re-parse).""" + x, y = self.wcs.world_to_pixel_values(float(ra_deg), float(dec_deg)) + return float(x), float(y) + + +def frame_geometry( + header: Mapping[str, Any], + aperture_radius_arcsec: float, + annulus_inner_radius_arcsec: float, + annulus_outer_radius_arcsec: float, +) -> FrameGeometry: + """ + Builds the reusable per-frame geometry: one WCS plus the three aperture radii converted to + pixels via the frame's plate scale. Matches arcsec_to_pixels/world_to_pixel exactly, just + without rebuilding the WCS for every candidate. + """ + pixel_scale = pixel_scale_arcsec(header) + return FrameGeometry( + wcs=WCS(dict(header)), + aperture_radius_px=float(aperture_radius_arcsec) / pixel_scale, + annulus_inner_radius_px=float(annulus_inner_radius_arcsec) / pixel_scale, + annulus_outer_radius_px=float(annulus_outer_radius_arcsec) / pixel_scale, + ) From 8c0402f8ed59cfd72f369e44ae687ba644e0f0f0 Mon Sep 17 00:00:00 2001 From: Jon Date: Fri, 17 Jul 2026 01:00:51 +0000 Subject: [PATCH 3/5] Update diagnostic images to be cropped and use auto zscaling and be stored as a jpeg --- .../data_operations/aperture_photometry.py | 22 +- .../tests/test_aperture_photometry.py | 72 +++--- .../datalab_session/tests/test_operations.py | 17 +- .../utils/aperture_light_curve.py | 56 +++-- .../datalab_session/utils/fits_metadata.py | 9 +- .../utils/photometry_diagnostics.py | 221 ++++++++---------- 6 files changed, 206 insertions(+), 191 deletions(-) diff --git a/datalab/datalab_session/data_operations/aperture_photometry.py b/datalab/datalab_session/data_operations/aperture_photometry.py index 4b7d5e4..7465bcf 100644 --- a/datalab/datalab_session/data_operations/aperture_photometry.py +++ b/datalab/datalab_session/data_operations/aperture_photometry.py @@ -6,7 +6,9 @@ from datalab.datalab_session.data_operations.data_operation import BaseDataOperation from datalab.datalab_session.exceptions import ClientAlertException from datalab.datalab_session.utils.filecache import FileCache +from datalab.datalab_session.utils.file_utils import temp_file_manager from datalab.datalab_session.utils.format import Format +from datalab.datalab_session.utils.s3_utils import save_files_to_s3 from datalab.datalab_session.utils.aperture_light_curve import ( DEFAULT_ANNULUS_INNER_RADIUS, DEFAULT_ANNULUS_OUTER_RADIUS, @@ -152,6 +154,7 @@ def operate(self, submitter: User): raise ClientAlertException(f'Operation {self.name()} received invalid input.') from exc self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['APERTURE_PHOTOMETRY_PERCENTAGE_COMPLETION']) + diagnostic_image_urls = self._save_diagnostic_images_to_s3(result.diagnostic_image_jpegs_by_fits_basename) filter_value = input_files[0].get('filter', input_files[0].get('primary_optical_element', 'None')) output = { 'output_data': [ @@ -166,7 +169,7 @@ def operate(self, submitter: User): asdict(star) for star in result.selected_comparison_stars ], 'diagnostics': result.diagnostics_by_fits_basename, - 'diagnostic_images': result.diagnostic_images_by_fits_basename, + 'diagnostic_images': diagnostic_image_urls, } ] } @@ -177,5 +180,20 @@ def operate(self, submitter: User): "Aperture Photometry output: " f"filter={filter_value}, light_curve_rows={len(result.light_curve_rows)}, " f"selected_comparison_stars={len(result.selected_comparison_stars)}, " - f"diagnostic_images={len(result.diagnostic_images_by_fits_basename)}" + f"diagnostic_images={len(diagnostic_image_urls)}" ) + + def _save_diagnostic_images_to_s3(self, diagnostic_image_jpegs_by_fits_basename: dict) -> dict: + """ + Uploads each frame's diagnostic overlay JPEG to the operation bucket. + + Returns a dict mapping FITS basename to the presigned bucket url for its overlay. + """ + diagnostic_image_urls = {} + for index, (fits_basename, jpeg_bytes) in enumerate(diagnostic_image_jpegs_by_fits_basename.items(), start=1): + with temp_file_manager(f'{self.cache_key}-{index}-diagnostic.jpg', dir=self.temp) as jpeg_path: + with open(jpeg_path, 'wb') as jpeg_file: + jpeg_file.write(jpeg_bytes) + s3_output = save_files_to_s3(self.cache_key, Format.IMAGE, {'diagnostic_jpg_path': jpeg_path}, index=index) + diagnostic_image_urls[fits_basename] = s3_output['diagnostic_url'] + return diagnostic_image_urls diff --git a/datalab/datalab_session/tests/test_aperture_photometry.py b/datalab/datalab_session/tests/test_aperture_photometry.py index 5436665..eef2ec4 100644 --- a/datalab/datalab_session/tests/test_aperture_photometry.py +++ b/datalab/datalab_session/tests/test_aperture_photometry.py @@ -4,7 +4,6 @@ import os import tempfile import unittest -import base64 from io import BytesIO from dataclasses import asdict from pathlib import Path @@ -568,20 +567,22 @@ def test_zero_point_guard_drops_blended_catalog_magnitude_outlier(self) -> None: self.assertNotIn("cand-blended", kept_ids) self.assertEqual(len(kept), 5) - def test_diagnostic_images_include_labeled_candidate_overlay_jpegs(self) -> None: + def test_diagnostic_images_include_candidate_overlay_jpegs(self) -> None: frames, (target_ra, target_dec) = build_frame_set() fits_paths = self.write_frames(frames) result = generate_light_curve(fits_paths, target_ra, target_dec, 4.0, 6.0, 9.0) self.assertEqual( - set(result.diagnostic_images_by_fits_basename), + set(result.diagnostic_image_jpegs_by_fits_basename), {"frame_1.fits", "frame_2.fits", "frame_3.fits"}, ) - image_data = base64.b64decode(result.diagnostic_images_by_fits_basename["frame_1.fits"]) + from datalab.datalab_session.utils.photometry_diagnostics import OVERLAY_MAX_DIMENSION + + image_data = result.diagnostic_image_jpegs_by_fits_basename["frame_1.fits"] image = Image.open(BytesIO(image_data)) self.assertEqual(image.format, "JPEG") - self.assertEqual(image.size, (80, 80)) + self.assertEqual(max(image.size), OVERLAY_MAX_DIMENSION) rgb = np.asarray(image.convert("RGB")) blue_overlay_pixels = np.count_nonzero( (rgb[:, :, 0] < 80) & @@ -602,7 +603,7 @@ def test_diagnostic_images_include_labeled_candidate_overlay_jpegs(self) -> None (rgb[:, :, 1] < 170) & (rgb[:, :, 2] < 90) )[:, 0] - self.assertGreater(float(np.mean(orange_rows)), 40.0) + self.assertGreater(float(np.mean(orange_rows)), rgb.shape[0] * 0.5) def test_determinism(self) -> None: frames, (target_ra, target_dec) = build_frame_set() @@ -620,9 +621,10 @@ def test_determinism(self) -> None: ) def test_pixel_data_streams_one_frame_at_a_time(self) -> None: - # The pipeline exists to keep memory flat in the input count: each frame's pixels must be - # loaded exactly once, at most one frame's pixels may be alive at any moment, and nothing - # in the result may retain pixel arrays after the run. + # The pipeline exists to keep memory flat in the input count: each frame's pixels are + # loaded once for measurement and once for overlay rendering, at most one frame's pixels + # may be alive at any moment, and nothing in the result may retain pixel arrays after + # the run. import gc import weakref from unittest import mock @@ -651,34 +653,14 @@ def tracking_load(fits_path: str): gc.collect() self.assertEqual(len(result.light_curve_rows), 3) - self.assertEqual(sorted(loaded_paths), sorted(fits_paths)) + self.assertEqual(sorted(loaded_paths), sorted(fits_paths * 2)) self.assertEqual(max_concurrent_images, 1) self.assertTrue(all(ref() is None for ref in loaded_refs)) - def test_frame_preview_downsamples_large_frames(self) -> None: - from datalab.datalab_session.utils.photometry_diagnostics import build_frame_preview - - image = np.full((2400, 3200), 100.0, dtype=np.float32) - - preview = build_frame_preview(image, max_dimension=1500) - - self.assertEqual(preview.gray.dtype, np.uint8) - self.assertEqual(preview.gray.shape, (800, 1066)) - self.assertLessEqual(max(preview.gray.shape), 1500) - self.assertAlmostEqual(preview.scale, 1.0 / 3.0) - - def test_frame_preview_keeps_small_frames_full_resolution(self) -> None: - from datalab.datalab_session.utils.photometry_diagnostics import build_frame_preview - - preview = build_frame_preview(np.full((80, 100), 100.0, dtype=np.float32)) - - self.assertEqual(preview.gray.shape, (80, 100)) - self.assertEqual(preview.scale, 1.0) - - def test_overlay_positions_scale_to_downsampled_preview(self) -> None: + def test_overlay_crops_full_resolution_before_resampling(self) -> None: from datalab.datalab_session.utils.photometry_diagnostics import ( - build_frame_preview, - candidate_overlay_jpeg_base64, + OVERLAY_MAX_DIMENSION, + candidate_overlay_jpeg_bytes, ) height, width = 2400, 3200 @@ -697,21 +679,25 @@ def test_overlay_positions_scale_to_downsampled_preview(self) -> None: "CD2_2": TEST_DEG_PER_PIXEL, } frame = SimpleNamespace(fits_path="big.fits", header=header, width=width, height=height) - preview = build_frame_preview(np.full((height, width), 100.0, dtype=np.float32)) - self.assertLess(preview.scale, 1.0) + image_data = np.full((height, width), 100.0, dtype=np.float32) + # Bright marker centered on the target's full-resolution position: the crop happens at + # full resolution around the target circle, and the circle must surround the marker. + image_data[297:304, 597:604] = 5000.0 target_full_res = SimpleNamespace(x=600.0, y=300.0) - encoded = candidate_overlay_jpeg_base64( + jpeg_bytes = candidate_overlay_jpeg_bytes( frame=frame, - preview=preview, + image=image_data, stars=[], measurements=[], target_measurement=target_full_res, aperture_radius=4.0, ) - rgb = np.asarray(Image.open(BytesIO(base64.b64decode(encoded))).convert("RGB")) - self.assertEqual(rgb.shape[:2], preview.gray.shape) + rgb = np.asarray(Image.open(BytesIO(jpeg_bytes)).convert("RGB")) + # The tight square crop around the single target circle is resampled up to the output + # size; cropping after the whole-frame downsample would have produced a 750x1000 image. + self.assertEqual(rgb.shape[:2], (OVERLAY_MAX_DIMENSION, OVERLAY_MAX_DIMENSION)) orange = np.argwhere( (rgb[:, :, 0] > 180) & (rgb[:, :, 1] > 80) & @@ -719,10 +705,10 @@ def test_overlay_positions_scale_to_downsampled_preview(self) -> None: (rgb[:, :, 2] < 90) ) self.assertGreater(len(orange), 0) - expected_x = 600.0 * preview.scale - expected_y = (preview.gray.shape[0] - 1) - 300.0 * preview.scale - self.assertAlmostEqual(float(np.mean(orange[:, 1])), expected_x, delta=3.0) - self.assertAlmostEqual(float(np.mean(orange[:, 0])), expected_y, delta=3.0) + marker = np.argwhere(np.all(rgb > 200, axis=2)) + self.assertGreater(len(marker), 0) + self.assertAlmostEqual(float(np.mean(orange[:, 1])), float(np.mean(marker[:, 1])), delta=5.0) + self.assertAlmostEqual(float(np.mean(orange[:, 0])), float(np.mean(marker[:, 0])), delta=5.0) def test_default_fits_dependencies_read_sci_header_and_cat_rows(self) -> None: frames, (target_ra, target_dec) = build_frame_set(frame_count=1) diff --git a/datalab/datalab_session/tests/test_operations.py b/datalab/datalab_session/tests/test_operations.py index 0cd63fd..524e3f8 100644 --- a/datalab/datalab_session/tests/test_operations.py +++ b/datalab/datalab_session/tests/test_operations.py @@ -1,6 +1,7 @@ from contextlib import contextmanager from datetime import datetime, timezone import shutil +import tempfile from types import SimpleNamespace from unittest import mock import math @@ -288,6 +289,7 @@ def valid_input_data(self): 'annulus_outer_radius': 19.10, } + @mock.patch('datalab.datalab_session.data_operations.aperture_photometry.save_files_to_s3') @mock.patch('datalab.datalab_session.data_operations.aperture_photometry.generate_light_curve') @mock.patch('datalab.datalab_session.data_operations.aperture_photometry.FileCache') @mock.patch.object(AperturePhotometry, 'set_status') @@ -300,8 +302,10 @@ def test_operate_requires_and_passes_explicit_radii_and_filter( mock_set_status, mock_file_cache, mock_generate_light_curve, + mock_save_files_to_s3, ): mock_file_cache.return_value.get_fits.return_value = '/tmp/fits_1.fits' + mock_save_files_to_s3.return_value = {'diagnostic_url': 'https://bucket/fits_1-diagnostic.jpg'} mock_generate_light_curve.return_value = SimpleNamespace( light_curve_rows=[ LightCurveRow( @@ -324,11 +328,13 @@ def test_operate_requires_and_passes_explicit_radii_and_filter( diagnostics_by_fits_basename={ 'fits_1.fits': ['loaded 1 frame', 'selected 5 comparison stars'], }, - diagnostic_images_by_fits_basename={}, + diagnostic_image_jpegs_by_fits_basename={'fits_1.fits': b'jpeg-bytes'}, ) input_data = self.valid_input_data() aperture_photometry = AperturePhotometry(input_data) + aperture_photometry.temp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, aperture_photometry.temp, ignore_errors=True) aperture_photometry.operate(None) mock_file_cache.return_value.get_fits.assert_called_once_with('fits_1', 'local', None) @@ -355,6 +361,13 @@ def test_operate_requires_and_passes_explicit_radii_and_filter( self.assertTrue(math.isnan( output['output_data'][0]['light_curve'][0]['target_calibrated_apparent_magnitude_uncertainty'] )) + mock_save_files_to_s3.assert_called_once_with( + aperture_photometry.cache_key, Format.IMAGE, mock.ANY, index=1 + ) + self.assertEqual( + output['output_data'][0]['diagnostic_images'], + {'fits_1.fits': 'https://bucket/fits_1-diagnostic.jpg'}, + ) mock_set_status.assert_called_once_with('COMPLETED') def test_operate_requires_aperture_radius(self): @@ -393,7 +406,7 @@ def test_operate_allows_missing_filter(self): selected_comparison_stars=[], diagnostics=[], diagnostics_by_fits_basename={}, - diagnostic_images_by_fits_basename={}, + diagnostic_image_jpegs_by_fits_basename={}, ) AperturePhotometry(input_data).operate(None) diff --git a/datalab/datalab_session/utils/aperture_light_curve.py b/datalab/datalab_session/utils/aperture_light_curve.py index f73ca5b..368c06b 100644 --- a/datalab/datalab_session/utils/aperture_light_curve.py +++ b/datalab/datalab_session/utils/aperture_light_curve.py @@ -34,9 +34,7 @@ minimum_angular_neighbor_distance_arcsec, ) from datalab.datalab_session.utils.photometry_diagnostics import ( - FramePreview, - build_frame_preview, - candidate_overlay_jpeg_base64, + candidate_overlay_jpeg_bytes, comparison_star_validation_diagnostics, ) from datalab.datalab_session.utils.photometry import measure_aperture @@ -154,7 +152,7 @@ class LightCurveResult: light_curve_rows: list[LightCurveRow] diagnostics: list[str] diagnostics_by_fits_basename: dict[str, list[str]] - diagnostic_images_by_fits_basename: dict[str, str] + diagnostic_image_jpegs_by_fits_basename: dict[str, bytes] def generate_light_curve( @@ -216,13 +214,12 @@ def generate_light_curve( candidate_stars = candidate_stars_from_catalog(catalog) target_measurements: dict[str, TargetMeasurement] = {} - previews: dict[str, FramePreview] = {} measurements_by_candidate: dict[str, dict[str, ComparisonMeasurement]] = { candidate.candidate_id: {} for candidate in candidate_stars } failed_candidate_ids: set[str] = set() for frame in frames: - target, frame_measurements, newly_failed, preview = _measure_frame_pixels( + target, frame_measurements, newly_failed = _measure_frame_pixels( frame=frame, candidate_stars=candidate_stars, skip_candidate_ids=failed_candidate_ids, @@ -233,7 +230,6 @@ def generate_light_curve( annulus_outer_radius=annulus_outer_radius, ) target_measurements[frame.fits_path] = target - previews[frame.fits_path] = preview failed_candidate_ids |= newly_failed for candidate_id in newly_failed: measurements_by_candidate.pop(candidate_id, None) @@ -280,7 +276,7 @@ def generate_light_curve( frame_results: list[FrameResult] = [] light_curve_rows: list[LightCurveRow] = [] - diagnostic_images_by_fits_basename: dict[str, str] = {} + diagnostic_image_jpegs_by_fits_basename: dict[str, bytes] = {} for frame in frames: target = target_measurements[frame.fits_path] # Reuse the per-frame measurements captured during selection rather than re-measuring the @@ -347,9 +343,8 @@ def generate_light_curve( ) diagnostics.extend(frame_diagnostics) diagnostics_by_fits_basename[os.path.basename(frame.fits_path)].extend(frame_diagnostics) - diagnostic_images_by_fits_basename[os.path.basename(frame.fits_path)] = candidate_overlay_jpeg_base64( + diagnostic_image_jpegs_by_fits_basename[os.path.basename(frame.fits_path)] = _render_frame_overlay( frame=frame, - preview=previews[frame.fits_path], stars=selection.selected_stars, measurements=comparison_measurements, target_measurement=target, @@ -392,7 +387,7 @@ def generate_light_curve( light_curve_rows=light_curve_rows, diagnostics=diagnostics, diagnostics_by_fits_basename=diagnostics_by_fits_basename, - diagnostic_images_by_fits_basename=diagnostic_images_by_fits_basename, + diagnostic_image_jpegs_by_fits_basename=diagnostic_image_jpegs_by_fits_basename, ) @@ -513,16 +508,16 @@ def _measure_frame_pixels( aperture_radius: float, annulus_inner_radius: float, annulus_outer_radius: float, -) -> tuple[TargetMeasurement, dict[str, ComparisonMeasurement], set[str], FramePreview]: +) -> tuple[TargetMeasurement, dict[str, ComparisonMeasurement], set[str]]: """ - Runs all pixel-dependent work for one frame: the target measurement, a measurement of every - comparison candidate (minus skip_candidate_ids), and the downsampled diagnostic preview. + Runs all pixel-dependent work for one frame: the target measurement and a measurement of + every comparison candidate (minus skip_candidate_ids). The full-resolution image exists only inside this function, so it is released before the caller moves on to the next frame. - Returns the target measurement, this frame's candidate measurements by candidate_id, the - ids of candidates that failed to measure on this frame, and the preview. + Returns the target measurement, this frame's candidate measurements by candidate_id, and + the ids of candidates that failed to measure on this frame. """ image = _load_frame_image(frame.fits_path) # Build the frame's WCS and pixel-space aperture radii once, then reuse them for the target and @@ -551,8 +546,33 @@ def _measure_frame_pixels( ) except LightCurveError: failed_candidate_ids.add(candidate.candidate_id) - preview = build_frame_preview(image) - return target_measurement, candidate_measurements, failed_candidate_ids, preview + return target_measurement, candidate_measurements, failed_candidate_ids + + +def _render_frame_overlay( + *, + frame: FrameContext, + stars: Sequence[ComparisonStar], + measurements: Sequence[ComparisonMeasurement], + target_measurement: TargetMeasurement, + aperture_radius: float, +) -> bytes: + """ + Reloads one frame's pixels and renders its diagnostic overlay, cropped at full resolution + around the drawn circles before resampling. + + The full-resolution image exists only inside this function, so overlay rendering keeps + peak memory flat no matter how many frames are submitted. + """ + image = _load_frame_image(frame.fits_path) + return candidate_overlay_jpeg_bytes( + frame=frame, + image=image, + stars=stars, + measurements=measurements, + target_measurement=target_measurement, + aperture_radius=aperture_radius, + ) def _cat_rows(data: Any) -> list[dict[str, Any]]: diff --git a/datalab/datalab_session/utils/fits_metadata.py b/datalab/datalab_session/utils/fits_metadata.py index d1863ae..7de1464 100644 --- a/datalab/datalab_session/utils/fits_metadata.py +++ b/datalab/datalab_session/utils/fits_metadata.py @@ -1,11 +1,18 @@ import math +import warnings from dataclasses import dataclass from typing import Any, Mapping import numpy as np -from astropy.wcs import WCS +from astropy.wcs import WCS, FITSFixedWarning from astropy.wcs.utils import proj_plane_pixel_scales +# Archive headers store the observatory location as OBSGEO-X/Y/Z; wcslib normalizes them to +# OBSGEO-L/B/H on every WCS parse and reports the change as a FITSFixedWarning. The fix is +# purely informational and a WCS is built per frame all over the photometry pipeline, so +# silence the category process-wide. +warnings.filterwarnings('ignore', category=FITSFixedWarning) + def world_to_pixel(header: Mapping[str, Any], ra_deg: float, dec_deg: float) -> tuple[float, float]: x, y = WCS(dict(header)).world_to_pixel_values(float(ra_deg), float(dec_deg)) diff --git a/datalab/datalab_session/utils/photometry_diagnostics.py b/datalab/datalab_session/utils/photometry_diagnostics.py index ff8288e..4434499 100644 --- a/datalab/datalab_session/utils/photometry_diagnostics.py +++ b/datalab/datalab_session/utils/photometry_diagnostics.py @@ -1,127 +1,93 @@ from __future__ import annotations -import base64 import math -from dataclasses import dataclass from io import BytesIO from typing import Any, Sequence import numpy as np -from PIL import Image, ImageDraw, ImageFont +from fits2image.scaling import calc_zscale_min_max, extract_samples, linear_scale +from PIL import Image, ImageDraw from datalab.datalab_session.utils.fits_metadata import arcsec_to_pixels, optional_float from datalab.datalab_session.utils.flux_to_mag import flux_to_mag COMPARISON_STAR_COLOR = (0, 173, 239) TARGET_COLOR = (243, 131, 33) -# Diagnostic overlays are rendered on a downsampled preview no larger than this on either side. -PREVIEW_MAX_DIMENSION = 2000 +# Diagnostic overlays are resampled so their long side is this many pixels. +OVERLAY_MAX_DIMENSION = 1000 -@dataclass(frozen=True) -class FramePreview: - """ - Downsampled display rendering of a frame, captured while the frame's pixels were loaded. - - gray is uint8, display-oriented (y-flipped), block-mean downsampled so that preview - coordinates = full-resolution coordinates * scale. Diagnostic overlays are drawn on this - preview, so full-resolution pixels never have to be reloaded or retained for rendering. - """ - gray: np.ndarray - scale: float - - @property - def height(self) -> int: - return int(self.gray.shape[0]) - - @property - def width(self) -> int: - return int(self.gray.shape[1]) - - -def build_frame_preview(image: np.ndarray, max_dimension: int = PREVIEW_MAX_DIMENSION) -> FramePreview: - """ - Builds the downsampled uint8 display preview for a frame. - - Blocks are mean-combined (point sampling would drop stars smaller than the sampling step), - and the display stretch is computed on the downsampled array, so the only full-resolution - temporary is the trimmed copy the block reshape may make. - """ - height, width = image.shape - step = max(1, math.ceil(max(height, width) / max_dimension)) - if step > 1: - trimmed = image[: (height // step) * step, : (width // step) * step] - blocks = trimmed.reshape(height // step, step, width // step, step) - small = blocks.mean(axis=(1, 3), dtype=np.float64) - else: - small = np.asarray(image, dtype=float) - gray = np.flip(_stretch_to_uint8(small), axis=0) - return FramePreview(gray=np.ascontiguousarray(gray), scale=1.0 / step) - - -def candidate_overlay_jpeg_base64( +def candidate_overlay_jpeg_bytes( *, frame: Any, - preview: FramePreview, + image: np.ndarray, stars: Sequence[Any], measurements: Sequence[Any], target_measurement: Any, aperture_radius: float, -) -> str: - image = Image.fromarray(preview.gray).convert("RGB") - draw = ImageDraw.Draw(image) - font = _diagnostic_overlay_font(preview.width, preview.height) +) -> bytes: + """ + Renders the candidate-star overlay for one frame from its full-resolution pixels. + + The image is cropped at full resolution to the region containing every drawn circle plus + a one-radius margin, then resampled so its long side is OVERLAY_MAX_DIMENSION, so tight + star fields keep native detail instead of inheriting a whole-frame downsample. + """ + height, width = image.shape stars_by_id = {star.candidate_id: star for star in stars} - min_dimension = max(min(preview.width, preview.height), 1) - aperture_radius_px = arcsec_to_pixels(frame.header, aperture_radius) * preview.scale - radius = max(aperture_radius_px, min_dimension * 0.018, 14.0) - line_width = max(3, int(round(min_dimension * 0.004))) - label_padding = max(3, int(round(min_dimension * 0.004))) + comparison_positions: list[tuple[float, float]] = [] for measurement in measurements: if measurement.candidate_id not in stars_by_id: continue - x = float(measurement.x) * preview.scale - y = _display_y(float(measurement.y) * preview.scale, preview.height) - if not math.isfinite(x) or not math.isfinite(y): - continue + x = float(measurement.x) + y = float(measurement.y) + if math.isfinite(x) and math.isfinite(y): + comparison_positions.append((x, y)) + + target_x = float(target_measurement.x) + target_y = float(target_measurement.y) + target_position = (target_x, target_y) if math.isfinite(target_x) and math.isfinite(target_y) else None + positions = comparison_positions + ([target_position] if target_position else []) + + radius_full = max(arcsec_to_pixels(frame.header, aperture_radius), 14.0) + x0, y0, x1, y1 = _crop_bounds(positions, pad=2.0 * radius_full, width=width, height=height) + crop = image[y0:y1, x0:x1] + crop_height, crop_width = crop.shape + + scale = OVERLAY_MAX_DIMENSION / max(crop_width, crop_height) + out_width = max(int(round(crop_width * scale)), 1) + out_height = max(int(round(crop_height * scale)), 1) + # Resample the raw flux before stretching: stretching first would clip bright star cores and + # quantize faint ones to 8 bits before they are averaged, dimming and blurring both. This also + # keeps the stretch percentiles computed on display-resolution data. + resampled = Image.fromarray(np.ascontiguousarray(crop, dtype=np.float32)).resize( + (out_width, out_height), Image.Resampling.LANCZOS + ) + gray = np.flip(_stretch_to_uint8(np.asarray(resampled)), axis=0) + overlay = Image.fromarray(np.ascontiguousarray(gray)).convert("RGB") + draw = ImageDraw.Draw(overlay) + scale_x = out_width / crop_width + scale_y = out_height / crop_height + min_dimension = min(out_width, out_height) + radius = max(radius_full * scale, min_dimension * 0.035, 24.0) + line_width = max(3, int(round(min_dimension * 0.004))) - label = _overlay_label(measurement.candidate_id) - halo_bbox = (x - radius, y - radius, x + radius, y + radius) - draw.ellipse(halo_bbox, outline=(0, 0, 0), width=line_width + 2) - draw.ellipse(halo_bbox, outline=COMPARISON_STAR_COLOR, width=line_width) - - label_x = x + radius + label_padding - label_y = y - radius - label_padding - label_bbox = draw.textbbox((label_x, label_y), label, font=font) - label_width = label_bbox[2] - label_bbox[0] - label_height = label_bbox[3] - label_bbox[1] - if label_x + label_width + label_padding > preview.width: - label_x = max(x - radius - label_width - label_padding, 0) - if label_y < 0: - label_y = min(y + radius + label_padding, max(preview.height - label_height - label_padding, 0)) - draw.text( - (label_x, label_y), - label, - fill=COMPARISON_STAR_COLOR, - font=font, - ) + # The image is y-flipped for display, and pixel centers map through a resize as + # (coordinate + 0.5) * scale - 0.5. + for x, y in comparison_positions: + cx = (x - x0 + 0.5) * scale_x - 0.5 + cy = ((crop_height - 1) - (y - y0) + 0.5) * scale_y - 0.5 + draw.ellipse((cx - radius, cy - radius, cx + radius, cy + radius), outline=COMPARISON_STAR_COLOR, width=line_width) - target_x = float(target_measurement.x) * preview.scale - target_y = _display_y(float(target_measurement.y) * preview.scale, preview.height) - if math.isfinite(target_x) and math.isfinite(target_y): - target_bbox = ( - target_x - radius, - target_y - radius, - target_x + radius, - target_y + radius, - ) - draw.ellipse(target_bbox, outline=(0, 0, 0), width=line_width + 2) - draw.ellipse(target_bbox, outline=TARGET_COLOR, width=line_width) + if target_position is not None: + cx = (target_position[0] - x0 + 0.5) * scale_x - 0.5 + cy = ((crop_height - 1) - (target_position[1] - y0) + 0.5) * scale_y - 0.5 + draw.ellipse((cx - radius, cy - radius, cx + radius, cy + radius), outline=TARGET_COLOR, width=line_width) buffer = BytesIO() - image.save(buffer, format="JPEG", quality=90) - return base64.b64encode(buffer.getvalue()).decode("utf-8") + overlay.save(buffer, format="JPEG", quality=90) + return buffer.getvalue() def comparison_star_validation_diagnostics( @@ -158,40 +124,45 @@ def comparison_star_validation_diagnostics( def _stretch_to_uint8(image_data: np.ndarray) -> np.ndarray: - finite = np.asarray(image_data, dtype=float) - finite_values = finite[np.isfinite(finite)] - if finite_values.size: - zmin, zmax = np.percentile(finite_values, (1, 99.7)) + """ + Stretches image data to uint8 for display with the same fits2image zscale autoscaling + used for the other datalab JPEGs (see fits2image.scaling.auto_scale): black point at the + sample median, white point from a zscale fit of the sorted samples, and a 2.5 gamma. + """ + finite = np.nan_to_num(np.asarray(image_data, dtype=float)) + height, width = finite.shape + samples = extract_samples(finite, {'NAXIS1': width, 'NAXIS2': height}, nsamples=min(2000, finite.size)) + if samples.size >= 10: + zmin = float(np.median(samples)) + _, zmax, _ = calc_zscale_min_max(samples, contrast=0.1, iterations=1) + zmax = float(zmax) + if not math.isfinite(zmin) or not math.isfinite(zmax) or zmax <= zmin: + zmin, zmax = float(np.min(finite)), float(np.max(finite)) else: - zmin, zmax = 0.0, 1.0 - if not math.isfinite(float(zmin)) or not math.isfinite(float(zmax)) or zmax <= zmin: - zmin = float(np.nanmin(finite_values)) if finite_values.size else 0.0 - zmax = float(np.nanmax(finite_values)) if finite_values.size else 1.0 - if zmax <= zmin: - zmax = zmin + 1.0 - - scaled = np.clip((finite - zmin) / (zmax - zmin), 0.0, 1.0) - scaled = np.nan_to_num(scaled, nan=0.0, posinf=1.0, neginf=0.0) - return (scaled * 255.0).astype(np.uint8) - - -def _display_y(y: float, height: int) -> float: - return float(height - 1) - y + zmin, zmax = float(np.min(finite)), float(np.max(finite)) + return linear_scale(finite, zmin, zmax, max_val=255, gamma_adjust=2.5) -def _diagnostic_overlay_font(width: int, height: int) -> ImageFont.ImageFont: - font_size = max(48, min(360, int(round(min(width, height) * 0.09)))) - try: - return ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) - except OSError: - return ImageFont.load_default() - - -def _overlay_label(candidate_id: str) -> str: - for prefix in ("cand-", "comp-"): - if candidate_id.startswith(prefix): - return candidate_id[len(prefix):] - return candidate_id +def _crop_bounds( + positions: Sequence[tuple[float, float]], + pad: float, + width: int, + height: int, +) -> tuple[int, int, int, int]: + """ + Bounds of the region containing every position plus pad on all sides, clamped to the + image, in full-resolution pixel coordinates. Falls back to the whole image when there + are no positions or the clamped region is degenerate. + """ + if not positions: + return 0, 0, width, height + x0 = max(int(math.floor(min(x for x, _ in positions) - pad)), 0) + y0 = max(int(math.floor(min(y for _, y in positions) - pad)), 0) + x1 = min(int(math.ceil(max(x for x, _ in positions) + pad)) + 1, width) + y1 = min(int(math.ceil(max(y for _, y in positions) + pad)) + 1, height) + if x0 >= x1 or y0 >= y1: + return 0, 0, width, height + return x0, y0, x1, y1 def _flux_to_magnitude(flux: float, zero_point: float) -> float: From 428836330de9e8745e06e3d774ba22a88394bc5a Mon Sep 17 00:00:00 2001 From: Jon Date: Fri, 17 Jul 2026 05:32:02 +0000 Subject: [PATCH 4/5] A few low hanging fruit optimizations to halve the speed of aperture photometry again. There are still further optimizations that can be done, but vectorizing the centroiding stuff, or by using a spatial hash for the cluster matching to only compute cluster angular distance for close clusters. --- .../utils/aperture_light_curve.py | 6 +-- datalab/datalab_session/utils/centroiding.py | 44 +++++++------------ .../datalab_session/utils/comparison_stars.py | 6 +-- .../datalab_session/utils/fits_metadata.py | 15 ++++--- datalab/datalab_session/utils/photometry.py | 14 ++++++ 5 files changed, 46 insertions(+), 39 deletions(-) diff --git a/datalab/datalab_session/utils/aperture_light_curve.py b/datalab/datalab_session/utils/aperture_light_curve.py index 368c06b..d890eae 100644 --- a/datalab/datalab_session/utils/aperture_light_curve.py +++ b/datalab/datalab_session/utils/aperture_light_curve.py @@ -21,9 +21,7 @@ from datalab.datalab_session.utils.fits_metadata import ( FrameGeometry, arcsec_to_pixels, - frame_gain, frame_geometry, - frame_read_noise, optional_float, world_to_pixel, ) @@ -698,8 +696,8 @@ def _measure_target( y_center=y_center, aperture_radius_px=aperture_radius_px, background_model=background_model, - gain=frame_gain(frame.header), - read_noise=frame_read_noise(frame.header), + gain=geometry.gain, + read_noise=geometry.read_noise, dark=0.0, error_class=LightCurveError, ) diff --git a/datalab/datalab_session/utils/centroiding.py b/datalab/datalab_session/utils/centroiding.py index dcf6944..4a37197 100644 --- a/datalab/datalab_session/utils/centroiding.py +++ b/datalab/datalab_session/utils/centroiding.py @@ -140,19 +140,19 @@ def calculate_background_model( j2 = int(y_center + r_back2) annulus_pixels: list[tuple[float, float, float]] = [] + for j in range(j1, j2 + 1): + dj = j - y_center + HALF_PIXEL + for i in range(i1, i2 + 1): + di = i - x_center + HALF_PIXEL + radius2 = di * di + dj * dj + if r12 <= radius2 <= r22: + value = _pixel(image, i, j) + if not math.isnan(value): + annulus_pixels.append((di, dj, value)) + + back_mean = 0.0 + back2_mean = 0.0 if remove_background_stars: - for j in range(j1, j2 + 1): - dj = j - y_center + HALF_PIXEL - for i in range(i1, i2 + 1): - di = i - x_center + HALF_PIXEL - radius2 = di * di + dj * dj - if r12 <= radius2 <= r22: - value = _pixel(image, i, j) - if not math.isnan(value): - annulus_pixels.append((di, dj, value)) - - back_mean = 0.0 - back2_mean = 0.0 previous_back_mean = 0.0 for iteration in range(max_iterations): back_stdev = math.sqrt(max(0.0, back2_mean - back_mean * back_mean)) @@ -169,26 +169,16 @@ def calculate_background_model( if abs(previous_back_mean - back_mean) < tolerance: break previous_back_mean = back_mean - else: - back_mean = 0.0 - back2_mean = 0.0 back_stdev = math.sqrt(max(0.0, back2_mean - back_mean * back_mean)) lower = back_mean - 2.0 * back_stdev upper = back_mean + 2.0 * back_stdev - kept: list[tuple[float, float, float]] = [] - for j in range(j1, j2 + 1): - dj = j - y_center + HALF_PIXEL - for i in range(i1, i2 + 1): - di = i - x_center + HALF_PIXEL - radius2 = di * di + dj * dj - if r12 <= radius2 <= r22: - value = _pixel(image, i, j) - if math.isnan(value): - continue - if not remove_background_stars or (lower <= value <= upper): - kept.append((di, dj, value)) + kept = [ + (di, dj, value) + for di, dj, value in annulus_pixels + if not remove_background_stars or (lower <= value <= upper) + ] background = sum(value for _, _, value in kept) / len(kept) if kept else 0.0 plane = _fit_plane(kept) if use_plane_background else None diff --git a/datalab/datalab_session/utils/comparison_stars.py b/datalab/datalab_session/utils/comparison_stars.py index 70d256e..1edc750 100644 --- a/datalab/datalab_session/utils/comparison_stars.py +++ b/datalab/datalab_session/utils/comparison_stars.py @@ -7,7 +7,7 @@ import numpy as np from datalab.datalab_session.utils.centroiding import centroid -from datalab.datalab_session.utils.fits_metadata import FrameGeometry, frame_gain, frame_read_noise +from datalab.datalab_session.utils.fits_metadata import FrameGeometry from datalab.datalab_session.utils.photometry import measure_aperture @@ -181,8 +181,8 @@ def measure_candidate_on_frame( y_center=centroid_result.y, aperture_radius_px=aperture_radius_px, background_model=centroid_result.background_model, - gain=frame_gain(frame.header), - read_noise=frame_read_noise(frame.header), + gain=geometry.gain, + read_noise=geometry.read_noise, dark=0.0, error_class=error_class, ) diff --git a/datalab/datalab_session/utils/fits_metadata.py b/datalab/datalab_session/utils/fits_metadata.py index 7de1464..f5bb56e 100644 --- a/datalab/datalab_session/utils/fits_metadata.py +++ b/datalab/datalab_session/utils/fits_metadata.py @@ -77,8 +77,8 @@ def arcsec_to_pixels(header: Mapping[str, Any], angular_radius_arcsec: float) -> @dataclass(frozen=True) class FrameGeometry: """ - Per-frame WCS and pixel-space aperture geometry, built once and reused for the target and - every comparison candidate on the frame. + Per-frame WCS, pixel-space aperture geometry, and detector noise parameters, built once + and reused for the target and every comparison candidate on the frame. Constructing a WCS from a header costs on the order of ~10 ms; arcsec_to_pixels and world_to_pixel each built one per call, so measuring a frame's candidates rebuilt it @@ -89,6 +89,8 @@ class FrameGeometry: aperture_radius_px: float annulus_inner_radius_px: float annulus_outer_radius_px: float + gain: float + read_noise: float def world_to_pixel(self, ra_deg: float, dec_deg: float) -> tuple[float, float]: """Pixel coordinates of a sky position using the cached WCS (no header re-parse).""" @@ -103,9 +105,10 @@ def frame_geometry( annulus_outer_radius_arcsec: float, ) -> FrameGeometry: """ - Builds the reusable per-frame geometry: one WCS plus the three aperture radii converted to - pixels via the frame's plate scale. Matches arcsec_to_pixels/world_to_pixel exactly, just - without rebuilding the WCS for every candidate. + Builds the reusable per-frame geometry: one WCS, the three aperture radii converted to + pixels via the frame's plate scale, and the detector gain and read noise. Matches + arcsec_to_pixels/world_to_pixel/frame_gain/frame_read_noise exactly, just without + re-deriving any of them for every candidate. """ pixel_scale = pixel_scale_arcsec(header) return FrameGeometry( @@ -113,4 +116,6 @@ def frame_geometry( aperture_radius_px=float(aperture_radius_arcsec) / pixel_scale, annulus_inner_radius_px=float(annulus_inner_radius_arcsec) / pixel_scale, annulus_outer_radius_px=float(annulus_outer_radius_arcsec) / pixel_scale, + gain=frame_gain(header), + read_noise=frame_read_noise(header), ) diff --git a/datalab/datalab_session/utils/photometry.py b/datalab/datalab_session/utils/photometry.py index e568c29..e65b325 100644 --- a/datalab/datalab_session/utils/photometry.py +++ b/datalab/datalab_session/utils/photometry.py @@ -83,6 +83,9 @@ def measure_aperture( } +PIXEL_HALF_DIAGONAL = math.sqrt(2.0) / 2.0 + + def fractional_pixel_overlap( i: int, j: int, @@ -94,6 +97,17 @@ def fractional_pixel_overlap( """ Approximates the fractional overlap of a pixel at (i, j) with a circular aperture centered at (x_center, y_center) with the given radius. """ + # Every subsample lies within the pixel's half-diagonal of its center, so pixels entirely + # outside or inside the circle resolve without evaluating the substeps grid. Only the + # one-pixel ring straddling the aperture edge needs subsampling. + dx_center = i + 0.5 - x_center + dy_center = j + 0.5 - y_center + center_distance = math.hypot(dx_center, dy_center) + if center_distance >= radius + PIXEL_HALF_DIAGONAL: + return 0.0 + if center_distance <= radius - PIXEL_HALF_DIAGONAL: + return 1.0 + inside = 0 total = substeps * substeps for sy in range(substeps): From 0ba236205ddbcf4b93a21408ad6e8da84639d599 Mon Sep 17 00:00:00 2001 From: Jon Date: Mon, 20 Jul 2026 22:22:35 +0000 Subject: [PATCH 5/5] Add reporting of progress per step and per frame within the step --- .../data_operations/aperture_photometry.py | 42 +++++++++++---- .../data_operations/data_operation.py | 5 ++ .../tests/test_aperture_photometry.py | 23 +++++++++ .../datalab_session/tests/test_operations.py | 1 + .../utils/aperture_light_curve.py | 51 +++++++++++++++---- 5 files changed, 102 insertions(+), 20 deletions(-) diff --git a/datalab/datalab_session/data_operations/aperture_photometry.py b/datalab/datalab_session/data_operations/aperture_photometry.py index 7465bcf..0e0f0e4 100644 --- a/datalab/datalab_session/data_operations/aperture_photometry.py +++ b/datalab/datalab_session/data_operations/aperture_photometry.py @@ -3,7 +3,7 @@ from django.contrib.auth.models import User -from datalab.datalab_session.data_operations.data_operation import BaseDataOperation +from datalab.datalab_session.data_operations.data_operation import BaseDataOperation, ProgressStep from datalab.datalab_session.exceptions import ClientAlertException from datalab.datalab_session.utils.filecache import FileCache from datalab.datalab_session.utils.file_utils import temp_file_manager @@ -33,9 +33,13 @@ class AperturePhotometry(BaseDataOperation): MINIMUM_NUMBER_OF_INPUTS = 1 MAXIMUM_NUMBER_OF_INPUTS = 999 PROGRESS_STEPS = { - 'INPUT_PROCESSING_PERCENTAGE_COMPLETION': 0.2, - 'APERTURE_PHOTOMETRY_PERCENTAGE_COMPLETION': 0.9, - 'OUTPUT_PERCENTAGE_COMPLETION': 1.0 + 'downloading': ProgressStep('Downloading input frames', 0.25), + 'validate': ProgressStep('Validating input frames', 0.3), + 'catalog': ProgressStep('Building comparison star catalog', 0.45), + 'measure': ProgressStep('Measuring source and comparison stars', 0.6), + 'select': ProgressStep('Selecting comparison stars', 0.75), + 'render': ProgressStep('Creating diagnostic images', 0.9), + 'save': ProgressStep('Saving output images', 1.0) } @staticmethod @@ -129,14 +133,14 @@ def operate(self, submitter: User): annulus_outer_radius = float(self.input_data['annulus_outer_radius']) min_comparisons = int(self.input_data.get('min_comparisons', DEFAULT_MIN_COMPARISONS)) max_comparisons = int(self.input_data.get('max_comparisons', DEFAULT_MAX_COMPARISONS)) - self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['INPUT_PROCESSING_PERCENTAGE_COMPLETION']) # Resolve inputs to local file-cache paths only. Pixel data is loaded (and released) # frame by frame inside generate_light_curve, never held for all inputs at once. file_cache = FileCache() - fits_paths = [ - file_cache.get_fits(input_file['basename'], input_file.get('source'), submitter) - for input_file in input_files - ] + fits_paths = [] + for index, input_file in enumerate(input_files, start=1): + fits_paths.append(file_cache.get_fits(input_file['basename'], input_file.get('source'), submitter)) + self._report_pipeline_progress('downloading', index / len(input_files)) + result = generate_light_curve( fits_paths=fits_paths, target_ra_deg=target_ra, @@ -146,6 +150,7 @@ def operate(self, submitter: User): annulus_outer_radius=annulus_outer_radius, min_comparisons=min_comparisons, max_comparisons=max_comparisons, + progress_callback=self._report_pipeline_progress, ) except LightCurveError as exc: log.warning(f"Aperture Photometry failed: {exc}") @@ -153,7 +158,6 @@ def operate(self, submitter: User): except (KeyError, TypeError, ValueError) as exc: raise ClientAlertException(f'Operation {self.name()} received invalid input.') from exc - self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['APERTURE_PHOTOMETRY_PERCENTAGE_COMPLETION']) diagnostic_image_urls = self._save_diagnostic_images_to_s3(result.diagnostic_image_jpegs_by_fits_basename) filter_value = input_files[0].get('filter', input_files[0].get('primary_optical_element', 'None')) output = { @@ -174,7 +178,8 @@ def operate(self, submitter: User): ] } self.set_output(output, is_raw=True) - self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['OUTPUT_PERCENTAGE_COMPLETION']) + self.set_operation_progress(1.0) + self.set_message("") self.set_status('COMPLETED') log.info( "Aperture Photometry output: " @@ -183,6 +188,19 @@ def operate(self, submitter: User): f"diagnostic_images={len(diagnostic_image_urls)}" ) + def _report_pipeline_progress(self, phase: str, fraction: float): + """ + Advances this operation's overall progress and status message through the PROGRESS_STEPS + band identified by phase. The band runs from the previous step's progress to this step's, + filling as fraction goes 0 -> 1. + """ + steps = list(AperturePhotometry.PROGRESS_STEPS.values()) + band = list(AperturePhotometry.PROGRESS_STEPS).index(phase) + band_start = steps[band - 1].progress if band > 0 else 0.0 + step = steps[band] + self.set_operation_progress(band_start + (step.progress - band_start) * fraction) + self.set_message(f"{step.message}: {fraction * 100:.0f}%") + def _save_diagnostic_images_to_s3(self, diagnostic_image_jpegs_by_fits_basename: dict) -> dict: """ Uploads each frame's diagnostic overlay JPEG to the operation bucket. @@ -190,10 +208,12 @@ def _save_diagnostic_images_to_s3(self, diagnostic_image_jpegs_by_fits_basename: Returns a dict mapping FITS basename to the presigned bucket url for its overlay. """ diagnostic_image_urls = {} + total = len(diagnostic_image_jpegs_by_fits_basename) for index, (fits_basename, jpeg_bytes) in enumerate(diagnostic_image_jpegs_by_fits_basename.items(), start=1): with temp_file_manager(f'{self.cache_key}-{index}-diagnostic.jpg', dir=self.temp) as jpeg_path: with open(jpeg_path, 'wb') as jpeg_file: jpeg_file.write(jpeg_bytes) s3_output = save_files_to_s3(self.cache_key, Format.IMAGE, {'diagnostic_jpg_path': jpeg_path}, index=index) diagnostic_image_urls[fits_basename] = s3_output['diagnostic_url'] + self._report_pipeline_progress('save', index / total) return diagnostic_image_urls diff --git a/datalab/datalab_session/data_operations/data_operation.py b/datalab/datalab_session/data_operations/data_operation.py index d217454..a8cc5e2 100644 --- a/datalab/datalab_session/data_operations/data_operation.py +++ b/datalab/datalab_session/data_operations/data_operation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections import namedtuple import hashlib import json import os @@ -17,6 +18,10 @@ log = logging.getLogger() log.setLevel(logging.INFO) +# ProgressStep will be used to tie a human readable message to the amount of time the step takes for a step +ProgressStep = namedtuple('ProgressStep', ['message', 'progress']) + + class BaseDataOperation(ABC): def __init__(self, input_data: dict = None): diff --git a/datalab/datalab_session/tests/test_aperture_photometry.py b/datalab/datalab_session/tests/test_aperture_photometry.py index eef2ec4..2fafe2b 100644 --- a/datalab/datalab_session/tests/test_aperture_photometry.py +++ b/datalab/datalab_session/tests/test_aperture_photometry.py @@ -605,6 +605,29 @@ def test_diagnostic_images_include_candidate_overlay_jpegs(self) -> None: )[:, 0] self.assertGreater(float(np.mean(orange_rows)), rgb.shape[0] * 0.5) + def test_progress_callback_reports_monotonic_per_frame_progress(self) -> None: + from datalab.datalab_session.utils.aperture_light_curve import PROGRESS_PHASES + + frames, (target_ra, target_dec) = build_frame_set() + fits_paths = self.write_frames(frames) + reported: list[tuple[str, float]] = [] + + generate_light_curve( + fits_paths, target_ra, target_dec, 4.0, 6.0, 9.0, + progress_callback=lambda phase, fraction: reported.append((phase, fraction)), + ) + + phases = [phase for phase, _ in reported] + # Phases arrive in pipeline order and every phase reports. + self.assertEqual(sorted(set(phases), key=PROGRESS_PHASES.index), list(PROGRESS_PHASES)) + self.assertEqual(phases, sorted(phases, key=PROGRESS_PHASES.index)) + # The frame-iterating phases report an increasing fraction once per frame, ending at 1.0. + for phase in ("validate", "catalog", "measure", "render"): + fractions = [fraction for reported_phase, fraction in reported if reported_phase == phase] + self.assertEqual(len(fractions), len(fits_paths)) + self.assertEqual(fractions, sorted(fractions)) + self.assertAlmostEqual(fractions[-1], 1.0) + def test_determinism(self) -> None: frames, (target_ra, target_dec) = build_frame_set() fits_paths = self.write_frames(frames) diff --git a/datalab/datalab_session/tests/test_operations.py b/datalab/datalab_session/tests/test_operations.py index 524e3f8..1fe1d32 100644 --- a/datalab/datalab_session/tests/test_operations.py +++ b/datalab/datalab_session/tests/test_operations.py @@ -347,6 +347,7 @@ def test_operate_requires_and_passes_explicit_radii_and_filter( annulus_outer_radius=19.10, min_comparisons=5, max_comparisons=10, + progress_callback=mock.ANY, ) output, is_raw = mock_set_output.call_args.args[0], mock_set_output.call_args.kwargs['is_raw'] self.assertTrue(is_raw) diff --git a/datalab/datalab_session/utils/aperture_light_curve.py b/datalab/datalab_session/utils/aperture_light_curve.py index d890eae..22d1fd3 100644 --- a/datalab/datalab_session/utils/aperture_light_curve.py +++ b/datalab/datalab_session/utils/aperture_light_curve.py @@ -3,7 +3,7 @@ import os from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Iterable, Mapping, Sequence +from typing import Any, Callable, Iterable, Mapping, Sequence import numpy as np from astropy.io import fits @@ -72,6 +72,12 @@ # so requiring presence in all frames discards good stars and collapses the candidate pool. Selected # stars are still measured via WCS on all frames, so the ensemble stays consistent frame-to-frame. COMPARISON_FRAME_COVERAGE_FRACTION = 0.8 +# Pipeline phases reported to progress_callback, in execution order. +PROGRESS_PHASES = ("validate", "catalog", "measure", "select", "render") + +# Receives (phase, fraction): phase is one of PROGRESS_PHASES and fraction is the completed +# share of that phase, in [0, 1]. +ProgressCallback = Callable[[str, float], None] class LightCurveError(ValueError): @@ -162,6 +168,7 @@ def generate_light_curve( annulus_outer_radius: float, min_comparisons: int = 5, max_comparisons: int = 10, + progress_callback: ProgressCallback | None = None, ) -> LightCurveResult: """ Generates a calibrated target light curve from local input FITS files, using comparison @@ -172,6 +179,9 @@ def generate_light_curve( every candidate, selects a comparison ensemble, and produces calibrated light curve rows with diagnostics for the frontend. At most one frame's full-resolution pixels are in memory at any point, so memory does not grow with the number of input frames. + + progress_callback, if given, receives (phase, completed fraction of that phase) with + phases from PROGRESS_PHASES; the frame-iterating phases report once per frame. """ log.info( "Aperture Photometry pipeline starting: " @@ -190,8 +200,15 @@ def generate_light_curve( max_comparisons=max_comparisons, ) + def report_progress(phase: str, fraction: float) -> None: + if progress_callback is not None: + progress_callback(phase, min(max(fraction, 0.0), 1.0)) + diagnostics: list[str] = [] - frames = _validated_frame_contexts(fits_paths) + frames = _validated_frame_contexts( + fits_paths, + on_frame=lambda index, total: report_progress("validate", index / total), + ) diagnostics_by_fits_basename: dict[str, list[str]] = { os.path.basename(frame.fits_path): [] @@ -204,6 +221,7 @@ def generate_light_curve( target_dec_deg=target_dec_deg, aperture_radius=aperture_radius, annulus_outer_radius=annulus_outer_radius, + on_frame=lambda index, total: report_progress("catalog", index / total), ) log.info( "Aperture Photometry comparison catalog built: " @@ -216,7 +234,7 @@ def generate_light_curve( candidate.candidate_id: {} for candidate in candidate_stars } failed_candidate_ids: set[str] = set() - for frame in frames: + for frame_index, frame in enumerate(frames, start=1): target, frame_measurements, newly_failed = _measure_frame_pixels( frame=frame, candidate_stars=candidate_stars, @@ -239,6 +257,7 @@ def generate_light_curve( f"net_counts={target.net_source_counts:.6f}, uncertainty={target.source_uncertainty:.6f}, " f"background={target.mean_background_per_pixel:.6f}, peak={target.peak_pixel_value:.6f}" ) + report_progress("measure", frame_index / len(frames)) target_mag_proxy = _target_magnitude_proxy(target_measurements.values()) log.info(f"Aperture Photometry target magnitude proxy: {target_mag_proxy:.6f}") @@ -257,6 +276,7 @@ def generate_light_curve( f"candidate_ids={[star.candidate_id for star in selection.selected_stars]}, " f"selection_diagnostics={len(selection.diagnostics)}" ) + report_progress("select", 1.0) reference_magnitudes = np.asarray( [star.reference_magnitude for star in selection.selected_stars], @@ -275,7 +295,7 @@ def generate_light_curve( frame_results: list[FrameResult] = [] light_curve_rows: list[LightCurveRow] = [] diagnostic_image_jpegs_by_fits_basename: dict[str, bytes] = {} - for frame in frames: + for frame_index, frame in enumerate(frames, start=1): target = target_measurements[frame.fits_path] # Reuse the per-frame measurements captured during selection rather than re-measuring the # same apertures on the same frames. @@ -373,6 +393,7 @@ def generate_light_curve( target_calibrated_apparent_magnitude_uncertainty=calibrated_mag_sigma, ) ) + report_progress("render", frame_index / len(frames)) log.info( "Aperture Photometry pipeline completed: " @@ -410,16 +431,20 @@ def _validate_inputs( raise LightCurveError("min_comparisons and max_comparisons must be positive and min_comparisons <= max_comparisons.") -def _validated_frame_contexts(fits_paths: Sequence[str]) -> list[FrameContext]: +def _validated_frame_contexts( + fits_paths: Sequence[str], + on_frame: Callable[[int, int], None] | None = None, +) -> list[FrameContext]: """ Builds validated frame metadata for each input FITS path. Reads only the SCI header and the CAT table -- never SCI pixel data -- so validation memory and time stay flat regardless of frame count or sensor size. Frames that fail validation - are ignored with a warning. + are ignored with a warning. on_frame, if given, is called as on_frame(index, total) after + each input path is processed, including rejected ones. """ frames: list[FrameContext] = [] - for fits_path in fits_paths: + for path_index, fits_path in enumerate(fits_paths, start=1): log.info(f"Aperture Photometry validating FITS frame: {fits_path}") try: with fits.open(fits_path) as hdul: @@ -470,6 +495,8 @@ def _validated_frame_contexts(fits_paths: Sequence[str]) -> list[FrameContext]: "Aperture Photometry ignoring input frame after validation error: " f"frame={fits_path}, error={exc}" ) + if on_frame is not None: + on_frame(path_index, len(fits_paths)) if not frames: raise LightCurveError("Aperture photometry requires at least 1 valid input file.") @@ -735,12 +762,14 @@ def _build_field_star_catalog( target_dec_deg: float, aperture_radius: float, annulus_outer_radius: float, + on_frame: Callable[[int, int], None] | None = None, ) -> list[dict[str, Any]]: """ Builds comp star candidates from the source catalogs across valid frames. Returns candidates detected in at least COMPARISON_FRAME_COVERAGE_FRACTION of the frames that - are not too close to the target or the edge of the image. + are not too close to the target or the edge of the image. on_frame, if given, is called as + on_frame(index, total) after each frame's rows are cross-matched. """ clusters: list[dict[str, Any]] = [] target_pixels = { @@ -748,7 +777,7 @@ def _build_field_star_catalog( for frame in frames } - for frame in frames: + for frame_index, frame in enumerate(frames, start=1): rows: list[dict[str, Any]] = [] for raw_row in frame.second_hdu_rows: try: @@ -764,6 +793,8 @@ def _build_field_star_catalog( "rejected_too_close_to_target=0, " f"rejected_too_close_to_edge=0, clusters_so_far={len(clusters)}" ) + if on_frame is not None: + on_frame(frame_index, len(frames)) continue ra_values = np.asarray([row["ra_deg"] for row in rows], dtype=float) @@ -837,6 +868,8 @@ def _build_field_star_catalog( f"rejected_too_close_to_target={rejected_for_target}, " f"rejected_too_close_to_edge={rejected_for_edge}, clusters_so_far={len(clusters)}" ) + if on_frame is not None: + on_frame(frame_index, len(frames)) catalog: list[dict[str, Any]] = [] rejected_for_coverage = 0