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..0e0f0e4 100644 --- a/datalab/datalab_session/data_operations/aperture_photometry.py +++ b/datalab/datalab_session/data_operations/aperture_photometry.py @@ -3,10 +3,12 @@ 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.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 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, @@ -31,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 @@ -127,13 +133,16 @@ 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']) - input_handlers = [ - InputDataHandler(submitter, input_file['basename'], input_file.get('source')) - for input_file in input_files - ] + # 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 = [] + 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( - input_handlers=input_handlers, + fits_paths=fits_paths, target_ra_deg=target_ra, target_dec_deg=target_dec, aperture_radius=aperture_radius, @@ -141,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}") @@ -148,7 +158,7 @@ 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 = { 'output_data': [ @@ -163,16 +173,47 @@ 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, } ] } 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: " 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 _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. + + 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 0cd6814..2fafe2b 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 @@ -33,11 +32,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 +178,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 +191,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,8 +314,8 @@ 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.fits_metadata import world_to_pixel + from datalab.datalab_session.utils.aperture_light_curve import _load_frame_image, _measure_target + from datalab.datalab_session.utils.fits_metadata import frame_geometry, world_to_pixel frames, _ = build_frame_set() header = frames["frame_1.fits"]["header"] @@ -338,11 +326,10 @@ 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) @@ -580,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) & @@ -614,7 +603,30 @@ 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_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() @@ -631,6 +643,96 @@ 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 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 + + 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 * 2)) + self.assertEqual(max_concurrent_images, 1) + self.assertTrue(all(ref() is None for ref in loaded_refs)) + + def test_overlay_crops_full_resolution_before_resampling(self) -> None: + from datalab.datalab_session.utils.photometry_diagnostics import ( + OVERLAY_MAX_DIMENSION, + candidate_overlay_jpeg_bytes, + ) + + 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) + 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) + + jpeg_bytes = candidate_overlay_jpeg_bytes( + frame=frame, + image=image_data, + stars=[], + measurements=[], + target_measurement=target_full_res, + aperture_radius=4.0, + ) + + 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) & + (rgb[:, :, 1] < 170) & + (rgb[:, :, 2] < 90) + ) + self.assertGreater(len(orange), 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) handle = tempfile.NamedTemporaryFile(suffix=".fits", delete=False) @@ -659,9 +761,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 +800,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 +851,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..1fe1d32 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,8 +289,9 @@ 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.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 +300,12 @@ 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, + mock_save_files_to_s3, ): - 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_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( @@ -325,15 +328,18 @@ 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) 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, @@ -341,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) @@ -355,6 +362,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): @@ -383,17 +397,17 @@ 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=[], 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 845344b..22d1fd3 100644 --- a/datalab/datalab_session/utils/aperture_light_curve.py +++ b/datalab/datalab_session/utils/aperture_light_curve.py @@ -3,22 +3,25 @@ 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 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 from datalab.datalab_session.utils.fits_metadata import ( + FrameGeometry, arcsec_to_pixels, - frame_gain, - frame_read_noise, + frame_geometry, optional_float, world_to_pixel, ) @@ -29,7 +32,7 @@ minimum_angular_neighbor_distance_arcsec, ) from datalab.datalab_session.utils.photometry_diagnostics import ( - candidate_overlay_jpeg_base64, + candidate_overlay_jpeg_bytes, comparison_star_validation_diagnostics, ) from datalab.datalab_session.utils.photometry import measure_aperture @@ -41,6 +44,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 @@ -59,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): @@ -68,12 +87,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 @@ -134,11 +156,11 @@ 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( - input_handlers: list[Any], + fits_paths: list[str], target_ra_deg: float, target_dec_deg: float, aperture_radius: float, @@ -146,23 +168,31 @@ 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 input FITS files, using comparison stars from the source catalog. + Generates a calibrated target light curve from local 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. + 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. + + 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: " - 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, @@ -170,53 +200,72 @@ 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(input_handlers) + 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): [] 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, + on_frame=lambda index, total: report_progress("catalog", index / total), + ) + log.info( + "Aperture Photometry comparison catalog built: " + f"valid_candidates={len(catalog)}" + ) + candidate_stars = candidate_stars_from_catalog(catalog) + + target_measurements: dict[str, TargetMeasurement] = {} + measurements_by_candidate: dict[str, dict[str, ComparisonMeasurement]] = { + candidate.candidate_id: {} for candidate in candidate_stars + } + failed_candidate_ids: set[str] = set() + for frame_index, frame in enumerate(frames, start=1): + target, frame_measurements, newly_failed = _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 + 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}), " 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)) - 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, @@ -227,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], @@ -244,8 +294,8 @@ def generate_light_curve( frame_results: list[FrameResult] = [] light_curve_rows: list[LightCurveRow] = [] - diagnostic_images_by_fits_basename: dict[str, str] = {} - for frame in frames: + diagnostic_image_jpegs_by_fits_basename: dict[str, bytes] = {} + 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. @@ -311,7 +361,7 @@ 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, stars=selection.selected_stars, measurements=comparison_measurements, @@ -343,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: " @@ -355,21 +406,21 @@ 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, ) 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 +431,31 @@ 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], + 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. on_frame, if given, is called as on_frame(index, total) after + each input path is processed, including rejected ones. + """ frames: list[FrameContext] = [] - for input_handler in input_handlers: - fits_path = input_handler.fits_file + for path_index, fits_path in enumerate(fits_paths, start=1): 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 +465,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: @@ -432,6 +495,8 @@ def _validated_frame_contexts(input_handlers: Sequence[Any]) -> list[FrameContex "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.") @@ -444,10 +509,101 @@ 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]]: + """ + 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, 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 + # 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, + ) + 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, + geometry=geometry, + candidate=candidate, + error_class=LightCurveError, + ) + except LightCurveError: + failed_candidate_ids.add(candidate.candidate_id) + 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]]: 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,26 +644,28 @@ def _validate_second_hdu(rows: Sequence[Mapping[str, Any]], fits_path: str) -> N 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. + 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. 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( @@ -516,7 +674,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 +693,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,13 +718,13 @@ 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, 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, ) @@ -604,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 = { @@ -617,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: @@ -633,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) @@ -706,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 diff --git a/datalab/datalab_session/utils/centroiding.py b/datalab/datalab_session/utils/centroiding.py index 6036814..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 @@ -257,7 +247,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..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 arcsec_to_pixels, frame_gain, frame_read_noise, world_to_pixel +from datalab.datalab_session.utils.fits_metadata import FrameGeometry from datalab.datalab_session.utils.photometry import measure_aperture @@ -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,15 +142,18 @@ def _source_catalog_sort_key(candidate: ComparisonStar, target_mag_proxy: float) 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 aperture photometry at the refined position, estimating the background, summing the @@ -161,12 +161,12 @@ 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=frame.image, + image=image, x_click=x, y_click=y, radius=aperture_radius_px, @@ -176,13 +176,13 @@ 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, 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, ) @@ -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/fits_metadata.py b/datalab/datalab_session/utils/fits_metadata.py index 934cea5..f5bb56e 100644 --- a/datalab/datalab_session/utils/fits_metadata.py +++ b/datalab/datalab_session/utils/fits_metadata.py @@ -1,10 +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)) @@ -64,3 +72,50 @@ 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, 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 + 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 + 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).""" + 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, 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( + 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, + 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): diff --git a/datalab/datalab_session/utils/photometry_diagnostics.py b/datalab/datalab_session/utils/photometry_diagnostics.py index bbb3007..4434499 100644 --- a/datalab/datalab_session/utils/photometry_diagnostics.py +++ b/datalab/datalab_session/utils/photometry_diagnostics.py @@ -1,82 +1,93 @@ from __future__ import annotations -import base64 import math 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 resampled so their long side is this many pixels. +OVERLAY_MAX_DIMENSION = 1000 -def candidate_overlay_jpeg_base64( +def candidate_overlay_jpeg_bytes( *, frame: Any, + image: np.ndarray, stars: Sequence[Any], measurements: Sequence[Any], target_measurement: Any, aperture_radius: float, -) -> str: - image = _normalize_image_for_jpeg(frame.image) - draw = ImageDraw.Draw(image) - font = _diagnostic_overlay_font(frame.width, frame.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(frame.width, frame.height), 1) - aperture_radius_px = arcsec_to_pixels(frame.header, aperture_radius) - 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) - y = _display_y(float(measurement.y), frame.height) - if not math.isfinite(x) or not math.isfinite(y): - continue - - 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 > frame.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)) - draw.text( - (label_x, label_y), - label, - fill=COMPARISON_STAR_COLOR, - font=font, - ) + 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 = _display_y(float(target_measurement.y), frame.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) + 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))) + + # 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) + + 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( @@ -112,43 +123,46 @@ def comparison_star_validation_diagnostics( return diagnostics -def _normalize_image_for_jpeg(image_data: np.ndarray) -> Image.Image: - 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)) +def _stretch_to_uint8(image_data: np.ndarray) -> np.ndarray: + """ + 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) - gray = (scaled * 255.0).astype(np.uint8) - gray = np.flip(gray, axis=0) - return Image.fromarray(gray).convert("RGB") - - -def _display_y(y: float, height: int) -> float: - return float(height - 1) - y - - -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 + zmin, zmax = float(np.min(finite)), float(np.max(finite)) + return linear_scale(finite, zmin, zmax, max_val=255, gamma_adjust=2.5) + + +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: