diff --git a/datalab/datalab_session/analysis/variable_star.py b/datalab/datalab_session/analysis/variable_star.py deleted file mode 100644 index f688634..0000000 --- a/datalab/datalab_session/analysis/variable_star.py +++ /dev/null @@ -1,131 +0,0 @@ -import logging - -from django.contrib.auth.models import User -from astropy.timeseries import LombScargle -from astropy.time import Time -import numpy as np - -from datalab.datalab_session.utils.file_utils import get_hdu -from datalab.datalab_session.utils.filecache import FileCache -from datalab.datalab_session.utils.flux_to_mag import flux_to_mag - -log = logging.getLogger(__name__) -log.setLevel(logging.INFO) - -def variable_star(input: dict, user: User): - """ - Function to perform variable star analysis on a given image. - - input (dict): Input dictionary containing - target_coords(dict): ra,dec coordinates for target star - images(list): - image(dict): - basename(str): Name of the image file to be analyzed - id(str): Unique identifier for the image - observation_date(str): Date of the observation - """ - - coords = input.get("target_coords") - target_ra = coords.get("ra") - target_dec = coords.get("dec") - - light_curve = [] - excluded_images = [] - flux_fallback = False - - # Loop through each image's catalog and extract the target source's mag/magerr for the light curve - for image in input.get("images", []): - basename = image.get('basename') - - try: - file_path = FileCache().get_fits(basename, input.get('source', 'archive'), user) - cat_hdu = get_hdu(file_path, extension='CAT') - except Exception as e: - log.error(f"Error retrieving catalog for image {basename}: {e}") - excluded_images.append(basename) - continue - - target_source = find_target_source(cat_hdu, target_ra, target_dec) - if target_source is None: - log.info(f"No source found matching target coordinates: RA={target_ra}, DEC={target_dec} in image {basename}") - excluded_images.append(basename) - continue - - try: - mag = target_source['mag'] - print(f'mag: {mag}') - magerr = target_source['magerr'] - except KeyError as e: - # If mag or magerr is not present, fallback convert flux to mag - mag, magerr = flux_to_mag(target_source['flux'], target_source['fluxerr']) - print(f'flux mag: {mag}') - flux_fallback = True - except Exception as e: - log.warning(f"Invalid magnitude or magnitude error for target in image {basename}") - excluded_images.append(basename) - continue - - light_curve.append({ - 'mag': mag, - 'magerr': magerr, - 'julian_date': Time(image.get("observation_date")).jd, - 'observation_date': image.get("observation_date") - }) - print(f'light_curve: {light_curve}') - - try: - frequency, power, period, fap = calculate_period(light_curve) - except Exception as e: - log.error(f"Error calculating period: {e}") - period, fap = 0, 0 - - return { - 'target_coords': coords, - 'light_curve': light_curve, - 'flux_fallback': flux_fallback, - 'excluded_images': excluded_images, - 'period': period, - 'fap': fap, - 'frequency': frequency, - 'power': power - } - -def find_target_source(cat_hdu, target_ra, target_dec): - """ - Find the source in the catalog relative to the target coordinates. - """ - cat_data = cat_hdu.data - MATCH_PRECISION = 0.001 - - if 'ra' not in cat_data.names or 'dec' not in cat_data.names: - log.warning("CAT data does not have ra or dec names!") - return None - for source in cat_data: - target_ra = float(target_ra) - target_dec = float(target_dec) - - if abs(source['ra'] - target_ra) <= MATCH_PRECISION and abs(source['dec'] - target_dec) <= MATCH_PRECISION: - return source - - return None - -def calculate_period(light_curve): - """ - Use the astropy lomb scargle to perform the periodogram analysis on the light curve - """ - ls = LombScargle( - [lc['julian_date'] for lc in light_curve], - [lc['mag'] for lc in light_curve], - [lc['magerr'] for lc in light_curve] - ) - - frequency, power = ls.autopower() - - # Find the best frequency - best_frequency = frequency[np.argmax(power)] - period = 1 / best_frequency - - fap = ls.false_alarm_probability(power.max()) - - log.info(f"Best period found: {period} days with FAP: {fap}") - return frequency, power, period, fap diff --git a/datalab/datalab_session/data_operations/light_curve.py b/datalab/datalab_session/data_operations/light_curve.py new file mode 100644 index 0000000..598cdd3 --- /dev/null +++ b/datalab/datalab_session/data_operations/light_curve.py @@ -0,0 +1,169 @@ +import logging + +from astropy.time import Time +from django.contrib.auth.models import User + +from datalab.datalab_session.data_operations.data_operation import BaseDataOperation +from datalab.datalab_session.exceptions import ClientAlertException +from datalab.datalab_session.utils.format import Format +from datalab.datalab_session.utils.file_utils import get_hdu +from datalab.datalab_session.utils.filecache import FileCache +from datalab.datalab_session.utils.flux_to_mag import flux_to_mag + +log = logging.getLogger(__name__) +log.setLevel(logging.INFO) + + +class LightCurve(BaseDataOperation): + MINIMUM_NUMBER_OF_INPUTS = 1 + MAXIMUM_NUMBER_OF_INPUTS = 999 + PROGRESS_STEPS = { + "LIGHT_CURVE_DONE": 1.0, + } + + @staticmethod + def name(): + return "Light Curve" + + @staticmethod + def description(): + return "Builds a light curve for a source from calibrated catalog photometry across input images." + + @staticmethod + def wizard_description(): + return { + "name": LightCurve.name(), + "description": LightCurve.description(), + "category": "image", + "inputs": { + "source": { + "name": "Source Star", + "type": Format.SOURCE, + "description": "The source star to analyze", + "name_lookup": True, + }, + "input_files": { + "name": "Input Files", + "description": "The input files to pull photometry from", + "type": Format.FITS, + "single_filter": True, + "filter_options": ["rp", "ip", "gp", "zs"], + "minimum": LightCurve.MINIMUM_NUMBER_OF_INPUTS, + "maximum": LightCurve.MAXIMUM_NUMBER_OF_INPUTS, + }, + }, + } + + def operate(self, submitter: User): + input_files = self._validate_inputs( + input_key="input_files", + minimum_inputs=self.MINIMUM_NUMBER_OF_INPUTS, + ) + source = self.input_data.get("source") + if not source: + raise ClientAlertException(f"Operation {self.name()} requires a source.") + + output = light_curve( + self.input_data, + submitter, + allow_flux_fallback=False, + progress_callback=self.set_operation_progress, + ) + if not output["light_curve"]: + raise ClientAlertException("No extractable photometry data found for this source.") + + output["source"] = source + output["filter"] = input_files[0].get("filter", input_files[0].get("primary_optical_element", "None")) + + self.set_output({"output_data": [output]}, is_raw=True) + self.set_operation_progress(LightCurve.PROGRESS_STEPS["LIGHT_CURVE_DONE"]) + self.set_status("COMPLETED") + + +def light_curve(input: dict, user: User, allow_flux_fallback=True, progress_callback=None): + """ + Build a light curve from a target coordinate and a set of image catalogs. + """ + coords = input.get("target_coords") or input.get("source") + target_ra = coords.get("ra") + target_dec = coords.get("dec") + images = input.get("images") or input.get("input_files") or [] + + light_curve_data = [] + excluded_images = [] + flux_fallback = False + + for index, image in enumerate(images, start=1): + basename = image.get("basename") + image_source = image.get("source") + if image_source is None and isinstance(input.get("source"), str): + image_source = input.get("source") + if image_source is None: + image_source = "archive" + + try: + file_path = FileCache().get_fits(basename, image_source, user) + cat_hdu = get_hdu(file_path, extension="CAT") + except Exception as e: + log.error(f"Error retrieving catalog for image {basename}: {e}") + excluded_images.append(basename) + continue + + if progress_callback: + progress_callback(index / len(images)) + + target_source = find_target_source(cat_hdu, target_ra, target_dec) + if target_source is None: + log.info(f"No source found matching target coordinates: RA={target_ra}, DEC={target_dec} in image {basename}") + excluded_images.append(basename) + continue + + try: + mag = target_source["mag"] + magerr = target_source["magerr"] + except KeyError: + if not allow_flux_fallback: + log.warning(f"No calibrated magnitude or magnitude error for source in image {basename}") + excluded_images.append(basename) + continue + mag, magerr = flux_to_mag(target_source["flux"], target_source["fluxerr"]) + flux_fallback = True + except Exception: + log.warning(f"Invalid magnitude or magnitude error for target in image {basename}") + excluded_images.append(basename) + continue + + light_curve_data.append({ + "mag": mag, + "magerr": magerr, + "julian_date": Time(image.get("observation_date")).jd, + "observation_date": image.get("observation_date"), + }) + + return { + "target_coords": coords, + "light_curve": light_curve_data, + "flux_fallback": flux_fallback, + "excluded_images": excluded_images, + } + + +def find_target_source(cat_hdu, target_ra, target_dec): + """ + Find the source in the catalog relative to the target coordinates. + """ + cat_data = cat_hdu.data + match_precision = 0.001 + + if "ra" not in cat_data.names or "dec" not in cat_data.names: + log.warning("CAT data does not have ra or dec names!") + return None + + target_ra = float(target_ra) + target_dec = float(target_dec) + + for source in cat_data: + if abs(source["ra"] - target_ra) <= match_precision and abs(source["dec"] - target_dec) <= match_precision: + return source + + return None diff --git a/datalab/datalab_session/data_operations/variable_star.py b/datalab/datalab_session/data_operations/variable_star.py index 9ce915d..295008b 100644 --- a/datalab/datalab_session/data_operations/variable_star.py +++ b/datalab/datalab_session/data_operations/variable_star.py @@ -2,15 +2,12 @@ import numpy as np from astropy.timeseries import LombScargle -from astropy.time import Time -import numpy as np from django.contrib.auth.models import User from datalab.datalab_session.data_operations.data_operation import BaseDataOperation +from datalab.datalab_session.data_operations.light_curve import light_curve from datalab.datalab_session.exceptions import ClientAlertException from datalab.datalab_session.utils.format import Format -from datalab.datalab_session.utils.file_utils import get_hdu -from datalab.datalab_session.utils.filecache import FileCache log = logging.getLogger() log.setLevel(logging.INFO) @@ -71,50 +68,24 @@ def operate(self, submitter: User): comment= f'Datalab Variable Star analysis for target ra {target_ra}, dec {target_dec} on images: {", ".join([image["basename"] for image in input_files])}' log.info(comment) - light_curve = [] - excluded_images = [] - flux_fallback = False - - for index, image in enumerate(input_files, start = 1): - basename = image.get('basename') - filter = image.get('filter', image.get('primary_optical_element', 'None')) - try: - file_path = FileCache().get_fits(basename, image.get('source', 'archive'), submitter) - cat_hdu = get_hdu(file_path, extension='CAT') - except Exception as e: - log.error(f"Error retrieving catalog for image {basename}: {e}") - excluded_images.append(basename) - continue - - self.set_operation_progress(VariableStar.PROGRESS_STEPS['DOWNLOADING_INPUT_FILES'] * (index / len(input_files))) - - target_source = find_target_source(cat_hdu, target_ra, target_dec) - - if target_source is None: - log.info(f"No source found matching coordinates: RA={target_ra}, DEC={target_dec} in image {basename}") - excluded_images.append(basename) - continue - - try: - mag = target_source['mag'] - magerr = target_source['magerr'] - except Exception as e: - log.warning(f"No calibrated magnitude or magnitude error for source in image {basename}") - excluded_images.append(basename) - continue - - light_curve.append({ - 'mag': mag, - 'magerr': magerr, - 'julian_date': Time(image.get("observation_date")).jd, - 'observation_date': image.get("observation_date") - }) + light_curve_result = light_curve( + self.input_data, + submitter, + allow_flux_fallback=False, + progress_callback=lambda progress: self.set_operation_progress( + VariableStar.PROGRESS_STEPS['DOWNLOADING_INPUT_FILES'] * progress + ) + ) + light_curve_data = light_curve_result['light_curve'] + excluded_images = light_curve_result['excluded_images'] + flux_fallback = light_curve_result['flux_fallback'] + filter = input_files[0].get('filter', input_files[0].get('primary_optical_element', 'None')) - if len(light_curve) < VariableStar.MINIMUM_NUMBER_OF_INPUTS: - raise ClientAlertException(f'Only found extractable photometry data in {len(light_curve)} inputs - this is not enough to get a period. Try increasing your date range to get more input images and try again.') + if len(light_curve_data) < VariableStar.MINIMUM_NUMBER_OF_INPUTS: + raise ClientAlertException(f'Only found extractable photometry data in {len(light_curve_data)} inputs - this is not enough to get a period. Try increasing your date range to get more input images and try again.') try: - frequency, power, period, fap = calculate_period(light_curve) + frequency, power, period, fap = calculate_period(light_curve_data) except Exception as e: log.error(f"Error calculating period: {e}") period, fap = 0, 0 @@ -124,7 +95,7 @@ def operate(self, submitter: User): output = { 'output_data': [ { - 'light_curve': light_curve, + 'light_curve': light_curve_data, 'flux_fallback': flux_fallback, 'excluded_images': excluded_images, 'source': source, @@ -136,33 +107,12 @@ def operate(self, submitter: User): } ] } - log.info(f'Variable Star output found best period of {period} with false alarm prob. of {fap} using {len(light_curve)} / {len(input_files)} images') + log.info(f'Variable Star output found best period of {period} with false alarm prob. of {fap} using {len(light_curve_data)} / {len(input_files)} images') self.set_output(output, is_raw=True) self.set_operation_progress(VariableStar.PROGRESS_STEPS['OUTPUT_PERCENTAGE_COMPLETION']) self.set_status('COMPLETED') - -def find_target_source(cat_hdu, target_ra, target_dec): - """ - Find the source in the catalog relative to the target coordinates. - """ - cat_data = cat_hdu.data - MATCH_PRECISION = 0.001 - - if 'ra' not in cat_data.names or 'dec' not in cat_data.names: - log.warning("CAT data does not have ra or dec names!") - return None - for source in cat_data: - target_ra = float(target_ra) - target_dec = float(target_dec) - - if abs(source['ra'] - target_ra) <= MATCH_PRECISION and abs(source['dec'] - target_dec) <= MATCH_PRECISION: - return source - - return None - - def calculate_period(light_curve): """ Use the astropy lomb scargle to perform the periodogram analysis on the light curve diff --git a/datalab/datalab_session/tests/test_analysis.py b/datalab/datalab_session/tests/test_analysis.py index 5afd0e1..7c3501d 100644 --- a/datalab/datalab_session/tests/test_analysis.py +++ b/datalab/datalab_session/tests/test_analysis.py @@ -8,6 +8,9 @@ from numpy.testing import assert_almost_equal from datalab.datalab_session.analysis import centroiding, line_profile, source_catalog +from datalab.datalab_session.data_operations import light_curve as light_curve_module + + class TestAnalysis(TestCase): analysis_test_path = 'datalab/datalab_session/tests/test_files/analysis/' analysis_fits_1_path = f'datalab/datalab_session/tests/test_files/fits_1.fits.fz' @@ -70,6 +73,58 @@ def test_source_catalog(self, mock_file_cache): self.assertAlmostEqual(output[0]['mag'], expected_mag) self.assertAlmostEqual(output[0]['magerr'], expected_magerr) + @mock.patch('datalab.datalab_session.data_operations.light_curve.find_target_source') + @mock.patch('datalab.datalab_session.data_operations.light_curve.get_hdu') + @mock.patch('datalab.datalab_session.data_operations.light_curve.FileCache') + def test_light_curve_uses_catalog_magnitude(self, mock_file_cache, mock_get_hdu, mock_find_target_source): + mock_file_cache.return_value.get_fits.return_value = self.analysis_fits_1_path + mock_get_hdu.return_value = SimpleNamespace() + mock_find_target_source.return_value = { + 'mag': 15.2, + 'magerr': 0.03, + } + + output = light_curve_module.light_curve({ + 'target_coords': {'ra': 10.0, 'dec': 20.0}, + 'images': [{ + 'basename': 'fits_1', + 'observation_date': '2024-01-01T00:00:00', + }], + 'source': 'archive', + }, None) + + self.assertEqual(output['excluded_images'], []) + self.assertFalse(output['flux_fallback']) + self.assertEqual(len(output['light_curve']), 1) + self.assertEqual(output['light_curve'][0]['mag'], 15.2) + self.assertEqual(output['light_curve'][0]['magerr'], 0.03) + + @mock.patch('datalab.datalab_session.data_operations.light_curve.find_target_source') + @mock.patch('datalab.datalab_session.data_operations.light_curve.get_hdu') + @mock.patch('datalab.datalab_session.data_operations.light_curve.FileCache') + def test_light_curve_falls_back_to_flux(self, mock_file_cache, mock_get_hdu, mock_find_target_source): + mock_file_cache.return_value.get_fits.return_value = self.analysis_fits_1_path + mock_get_hdu.return_value = SimpleNamespace() + mock_find_target_source.return_value = { + 'flux': 100.0, + 'fluxerr': 5.0, + } + + output = light_curve_module.light_curve({ + 'target_coords': {'ra': 10.0, 'dec': 20.0}, + 'images': [{ + 'basename': 'fits_1', + 'observation_date': '2024-01-01T00:00:00', + }], + 'source': 'archive', + }, None) + + self.assertEqual(output['excluded_images'], []) + self.assertTrue(output['flux_fallback']) + self.assertEqual(len(output['light_curve']), 1) + self.assertAlmostEqual(output['light_curve'][0]['mag'], -5.0) + self.assertAlmostEqual(output['light_curve'][0]['magerr'], 0.05428681023790647) + def test_centroid_finds_pixels_center(self): image = np.zeros((21, 21), dtype=float) image[10, 10] = 100.0 diff --git a/datalab/datalab_session/tests/test_operations.py b/datalab/datalab_session/tests/test_operations.py index 6ccb497..99dddf9 100644 --- a/datalab/datalab_session/tests/test_operations.py +++ b/datalab/datalab_session/tests/test_operations.py @@ -9,6 +9,7 @@ from datalab.datalab_session.data_operations.data_operation import BaseDataOperation from datalab.datalab_session.data_operations.color_image import Color_Image from datalab.datalab_session.exceptions import ClientAlertException +from datalab.datalab_session.data_operations.light_curve import LightCurve from datalab.datalab_session.data_operations.median import Median from datalab.datalab_session.data_operations.stacking import Stack from datalab.datalab_session.tests.test_files.file_extended_test_case import FileExtendedTestCase @@ -222,6 +223,51 @@ def test_not_enough_files(self): median.operate(None) +class TestLightCurveOperation(FileExtendedTestCase): + + @mock.patch('datalab.datalab_session.data_operations.light_curve.light_curve') + def test_operate(self, mock_light_curve): + mock_light_curve.return_value = { + 'target_coords': {'ra': 10.0, 'dec': 20.0}, + 'light_curve': [{ + 'mag': 15.2, + 'magerr': 0.03, + 'julian_date': 2460310.5, + 'observation_date': '2024-01-01T00:00:00', + }], + 'flux_fallback': False, + 'excluded_images': [], + } + input_data = { + 'source': {'ra': 10.0, 'dec': 20.0}, + 'input_files': [{ + 'basename': 'fits_1', + 'source': 'local', + 'filter': 'rp', + 'observation_date': '2024-01-01T00:00:00', + }], + } + + light_curve = LightCurve(input_data) + light_curve.operate(None) + output = light_curve.get_output() + + self.assertEqual(light_curve.get_operation_progress(), 1.0) + self.assertEqual(light_curve.get_status(), 'COMPLETED') + self.assertEqual(output['output_data'][0]['source'], input_data['source']) + self.assertEqual(output['output_data'][0]['filter'], 'rp') + self.assertEqual(output['output_data'][0]['light_curve'], mock_light_curve.return_value['light_curve']) + + def test_not_enough_files(self): + light_curve = LightCurve({ + 'source': {'ra': 10.0, 'dec': 20.0}, + 'input_files': [], + }) + + with self.assertRaises(ClientAlertException): + light_curve.operate(None) + + class TestColorImageOperation(FileExtendedTestCase): temp_color_path = f'{test_path}temp_color.fits' test_color_path = f'{test_path}color_image/color_image.fits' diff --git a/datalab/datalab_session/views.py b/datalab/datalab_session/views.py index 35a4db7..12b3558 100644 --- a/datalab/datalab_session/views.py +++ b/datalab/datalab_session/views.py @@ -11,7 +11,6 @@ from datalab.datalab_session.analysis.get_tif import get_tif from datalab.datalab_session.analysis.get_jpg import get_jpg from datalab.datalab_session.analysis.raw_data import raw_data -from datalab.datalab_session.analysis.variable_star import variable_star from datalab.datalab_session.analysis.wcs import wcs from datalab.datalab_session.exceptions import ClientAlertException @@ -39,8 +38,7 @@ class AnalysisView(RetrieveAPIView): "get-tif": get_tif, "get-jpg": get_jpg, "raw-data": raw_data, - "wcs": wcs, - "variable-star": variable_star, + "wcs": wcs } def post(self, request, action):