From 5135238307f85a933f57161492c2d42d439f6d90 Mon Sep 17 00:00:00 2001 From: Kira Date: Wed, 1 Jul 2026 21:54:16 -0700 Subject: [PATCH] misc updates to script/notebook setup --- README.md | 10 +- .../Example_LSST_Catalog_Crossmatching.ipynb | 261 ++++++++---------- ...xample_Visualizing_LSST_Observations.ipynb | 14 +- src/thor/crossmatch_alerts.py | 223 ++++++++++++--- src/thor/utils/fetch_alerts.py | 83 +++++- 5 files changed, 377 insertions(+), 214 deletions(-) diff --git a/README.md b/README.md index ee64e44..47f2408 100644 --- a/README.md +++ b/README.md @@ -24,12 +24,20 @@ To fetch and visualise LSST alerts for a given night, run: This uses the Babamul alert broker to fetch alerts, which requires user credentials stored in a .env file locally. Omit dates to default to the previous night. The generated skymap is saved to `data/plots/` +To fetch alerts, crossmatch against available catalogs, and save candidates, run: + +```bash +python src/thor/crossmatch_alerts.py --start 06-28-2026 --end 06-30-2026 --additional_filtering tde_filter --scan +``` + +If there are any matches, the crossmatch result is saved to `data/lsst_alert_download/crossmatch_candidates_{timestamp}.json`. To apply additional TDE-specific filtering, pass `--additional_filtering tde_filter`. Add flag `--save` to save raw alerts to `data/lsst_alert_download/raw_files/`, and flag `--save_results` in order to save the crossmatch details locally to `data/lsst_alert_download/`. A summary of results will be printed in command line, but the --scan flag can also be included to open a temp jupyter notebook in browser and use Babamul's scanning tool. + ### Data Large data files (catalogs, simulation results, alert stream outputs) are not tracked in this repository. They are all saved locally in `THOR/data`. -In `data/catalogs` we keep all catalogs used for crossmatching. These catalogs are recorded in [`catalogs_catalog`](src/thor/alert_stream/catalogs_catalog.py). Most catalogs used have had basic quality cuts to select only galaxies, and have been reduced to selected columns. +In `data/catalogs` we keep all catalogs used for crossmatching. These catalogs are recorded in [`catalogs_catalog`](src/thor/catalogs_catalog.py). Most catalogs used have had basic quality cuts to select only galaxies, and have been reduced to selected columns. In `data/lsst_alert_download` we save LSST alerts fetched with the [Babamul alert broker](https://babamul.caltech.edu/). diff --git a/docs/notebooks/Example_LSST_Catalog_Crossmatching.ipynb b/docs/notebooks/Example_LSST_Catalog_Crossmatching.ipynb index a470f1b..db866e2 100644 --- a/docs/notebooks/Example_LSST_Catalog_Crossmatching.ipynb +++ b/docs/notebooks/Example_LSST_Catalog_Crossmatching.ipynb @@ -5,28 +5,15 @@ "id": "101160e4", "metadata": {}, "source": [ - "# fetch alerts with Babamul api and crossmatch / filter\n", + "# Filter LSST alerts to select high redshift objects\n", "\n", - "Adaped directly from: https://github.com/boom-astro/babamul/blob/main/examples/stream-basic/notebook.ipynb" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "719bfe87", - "metadata": {}, - "outputs": [], - "source": [ - "# TODOS\n", - "\n", - "# Make run daily, logging for which dates have been run, number filtered etc.\n", - "# create group on fritz and save objects there\n", - "# fully automate up to sending objects to fritz\n", - "# filter cosmos catalog based on recs here: https://cosmos2025.iap.fr/tutorials/catalog_usage_tutorial.html\n", - "# get other catalogs relevant for crossmatches with other fields (EDFS, M49)\n", - "# make filter version w/ fewer detections for anirudh (jet, grb afterglow))\n", - "# second pass to improve TDE filter, check other objects?\n", - "# redownload cosmos2025 and save more catalog values (ie z high and low, chi2, etc)" + "[Getting LSST alerts with the Babamul API](https://github.com/boom-astro/babamul/blob/main/examples/stream-basic/notebook.ipynb)\n", + "\n", + "### Notes on running this notebook:\n", + "\n", + "This notebook fetches alerts using the Babamul broker, which requires [user credentials saved in a .env file](https://github.com/boom-astro/babamul/blob/main/examples/stream-basic/.env.example). \n", + "\n", + "Babamul has periodically changed the naming schema, for example the token was previously changed from BABAMUL_API_TOKEN to BABAMUL_TOKEN, so the .env.example linked here may be incorrect." ] }, { @@ -37,28 +24,10 @@ "outputs": [], "source": [ "#imports \n", - "from astropy.table import Table\n", - "from datetime import datetime\n", - "import babamul\n", - "from babamul import jupyter\n", "import dotenv\n", "dotenv.load_dotenv()" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "3810bfa0", - "metadata": {}, - "outputs": [], - "source": [ - "import importlib\n", - "import rubin_stats_functions\n", - "import filter_functions\n", - "importlib.reload(rubin_stats_functions)\n", - "importlib.reload(filter_functions)" - ] - }, { "cell_type": "markdown", "id": "3558fb0f", @@ -66,9 +35,11 @@ "source": [ "### Catalog uniformization\n", "\n", - "cut to subsection of confident galaxies\n", + "The following code needs to be edited to work for each specific catalog.\n", + "\n", + "We cut the full catalog to subsection of confident galaxies\n", "\n", - "rename columns with universal schema (id, ra, dec, z, + any custom values that are useful)" + "We rename columns with standardized schema (id, ra, dec, z, and any unique values that we want to keep for a given catalog)" ] }, { @@ -80,7 +51,8 @@ "source": [ "from astropy.table import Table \n", "from astropy.io import fits\n", - "import numpy as np" + "import numpy as np\n", + "import matplotlib.pyplot as plt" ] }, { @@ -90,19 +62,13 @@ "metadata": {}, "outputs": [], "source": [ - "with fits.open('data/catalogs/COSMOS_11bands-SExtractor-Lephare.fits') as hdul:\n", - " hdul.info()\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0eba956a", - "metadata": {}, - "outputs": [], - "source": [ - "t = Table.read('data/catalogs/XMMLSS_11bands-SExtractor-Lephare.fits') \n", - "t[0:5]\n" + "# we load the full downloaded catalog\n", + "\n", + "with fits.open('../../data/catalogs/COSMOS_11bands-SExtractor-Lephare.fits') as hdul:\n", + " hdul.info()\n", + "\n", + "t = Table.read('../../data/catalogs/COSMOS_11bands-SExtractor-Lephare.fits') \n", + "t[0:5]" ] }, { @@ -112,7 +78,7 @@ "metadata": {}, "outputs": [], "source": [ - "# histogram t table ZPHOT values\n", + "# inspect catalog contents, ie histogram redshift values\n", "\n", "import matplotlib.pyplot as plt\n", "\n", @@ -126,6 +92,9 @@ "metadata": {}, "outputs": [], "source": [ + "# values to select on are specific to the catalog\n", + "\n", + "# here are our cuts for the clauds catalogs (https://www.clauds.net/available-data)\n", "galaxies = t[\n", " (t['OBJ_TYPE'] == 0) &\n", " (t['CHI_BEST'] > 0) & # has a valid SED fit\n", @@ -134,17 +103,10 @@ " (t['CLASS_STAR_HSC_I'] < 0.9) # morphologically not a star\n", "]\n", "\n", - "print(f'total objects: {len(t)}')\n", - "print(len(galaxies))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b48d0abe", - "metadata": {}, - "outputs": [], - "source": [ + "print(f'selected {len(galaxies)} from {len(t)} total objects')\n", + "\n", + "\n", + "# and then here is where we standardize column names from the Clauds names\n", "cut_catalog = galaxies['ID', 'RA', 'DEC', 'ZPHOT',\n", " 'Z_BEST68_LOW', 'Z_BEST68_HIGH', 'CHI_BEST',\n", " 'OBJ_TYPE', 'CLASS_STAR_HSC_I',\n", @@ -155,29 +117,14 @@ "cut_catalog.rename_column('DEC', 'dec')\n", "cut_catalog.rename_column('ZPHOT', 'z')\n", "\n", - "print(len(cut_catalog))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d6cb30aa", - "metadata": {}, - "outputs": [], - "source": [ + "# and then we print some statistics on the selected objects, which are copied into catalogs_catalog.py\n", + "\n", "for col in ['ra', 'dec', 'z']: \n", " print(f\"{col}: min={cut_catalog[col].min():.4f}, max={cut_catalog[col].max():.4f}\") \n", "\n", - "print(f'median z: {np.median(cut_catalog['z']):.4f}')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4c293afe", - "metadata": {}, - "outputs": [], - "source": [ + "print(f'median z: {np.median(cut_catalog['z']):.4f}')\n", + "\n", + "# finally we save the cut catalog to a fits file\n", "cut_catalog.write('data/catalogs/CLAUDS_SourceExtractor_XMMLSS_cut.fits', overwrite=False) " ] }, @@ -186,9 +133,9 @@ "id": "a8f9ee15", "metadata": {}, "source": [ - "### catalog crossmatch\n", + "### Fetching, prefiltering, organizing LSST alerts\n", "\n", - "Develop pipeline to crossmatch (or do any inference) on sources outside of Fritz, and save back to Fritz" + "Develop pipeline to crossmatch (or do any inference) on sources locally" ] }, { @@ -201,10 +148,10 @@ "# get alerts from date range and other basic filters for astrophysical sources\n", "# this has crashed for large date ranges, so is built to save in batches.\n", "\n", - "from rubin_stats_functions import babamul_get_alerts\n", + "from thor.utils.fetch_alerts import babamul_get_alerts\n", "\n", - "start=\"06-17-2026\"\n", - "end=\"06-29-2026\"\n", + "start=\"07-01-2026\"\n", + "end=\"07-02-2026\"\n", "\n", "loaded_alerts = babamul_get_alerts(\n", " survey=\"LSST\",\n", @@ -225,14 +172,15 @@ "metadata": {}, "outputs": [], "source": [ - "#if we saved batches as we fetched, combine the batched files of alerts into a single compressed file\n", - "from rubin_stats_functions import combine_alert_files\n", + "# Bookeeping: if we saved batches as we fetched, combine the batched files of alerts into a single compressed file and move out of raw_files\n", + "from thor.utils.fetch_alerts import combine_alert_files\n", + "from datetime import datetime\n", "\n", "start_fmt = datetime.strptime(start, \"%m-%d-%Y\").strftime(\"%m%d%y\") \n", "end_fmt = datetime.strptime(end, \"%m-%d-%Y\").strftime(\"%m%d%y\")\n", "filename = f\"basic_{start_fmt}_{end_fmt}.json.gz\" \n", "\n", - "combined = combine_alert_files(input_dir=\"data/lsst_alert_download/raw_files/\", output_path=f\"data/lsst_alert_download/{filename}\", delete_raw=True)" + "combined = combine_alert_files(input_dir=\"../../data/lsst_alert_download/raw_files/\", output_path=f\"../../data/lsst_alert_download/{filename}\", delete_raw=True)" ] }, { @@ -242,13 +190,13 @@ "metadata": {}, "outputs": [], "source": [ - "# bookkeeping - combine saved alerts to bigger batches per file\n", - "from rubin_stats_functions import combine_alert_files\n", + "# Bookkeeping: combine saved alerts to bigger batches per file\n", + "from thor.utils.fetch_alerts import combine_alert_files\n", "\n", "combined = combine_alert_files( \n", - " input_dir=\"data/lsst_alert_download/\", \n", + " input_dir=\"../../data/lsst_alert_download/\", \n", " input_files=[\"basic_041226_061726.json.gz\", \"basic_061726_062926.json.gz\"],\n", - " output_path=\"data/lsst_alert_download/basic_041226_062926.json.gz\",\n", + " output_path=\"../../data/lsst_alert_download/basic_041226_062926.json.gz\",\n", " delete_raw=True \n", ") " ] @@ -260,10 +208,10 @@ "metadata": {}, "outputs": [], "source": [ - "# if we want to open saved alerts file, we can do that too\n", + "# if we want to directly open a locally saved alerts file, we can do that too\n", "\n", - "from rubin_stats_functions import load_alerts\n", - "loaded_alerts = load_alerts(\"data/lsst_alert_download/basic_041226_062926.json.gz\")" + "from thor.utils.fetch_alerts import load_alerts\n", + "loaded_alerts = load_alerts(\"../../data/lsst_alert_download/basic_041226_062926.json.gz\")" ] }, { @@ -279,13 +227,15 @@ "# filter out dipole (image subtraction artifacts), poor PSF fits, low SNR, extended sources, and alerts with shape or centroid fitting issues. \n", "# filter currently repeats cuts made in babamul_get_alerts: drb<0.4, rock, star, near_brightstar, isdiffpos, and stationary.\n", "\n", - "filtered_alerts = filter_functions.filter_alerts( \n", + "from thor.utils.filter_functions import (filter_alerts, generic_filter, deduplicate_alerts)\n", + "\n", + "filtered_alerts = filter_alerts( \n", " loaded_alerts, \n", - " filter_functions.generic_filter, \n", + " generic_filter, \n", ") \n", "\n", "# deduplicate alerts to get unique objects\n", - "filtered_objects = filter_functions.deduplicate_alerts(filtered_alerts)" + "filtered_objects = deduplicate_alerts(filtered_alerts)" ] }, { @@ -296,20 +246,6 @@ "### Catalog crossmatch" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "9a44f4fa", - "metadata": {}, - "outputs": [], - "source": [ - "import importlib\n", - "import rubin_stats_functions\n", - "import filter_functions\n", - "importlib.reload(rubin_stats_functions)\n", - "importlib.reload(filter_functions)" - ] - }, { "cell_type": "code", "execution_count": null, @@ -317,14 +253,16 @@ "metadata": {}, "outputs": [], "source": [ - "# crossmatch with all available catalogs. returns a dictionary with all matched alert and catalog info.\n", + "# crossmatch with all available catalogs, ie all .fits files in data/catalogs\n", + "# returns a dictionary with all matched alert and catalog info.\n", + "# can provide catalog_name as a string or list of strings to specify which catalogs in catalog_path (default data/catalogs) to use. \n", + "# user must download catalogs and ensure there are columns with names id, ra, dec, and z - the function is not flexible in \n", + "# identifying different naming schema for these columns.\n", + "\n", + "from thor.utils.filter_functions import catalog_crossmatch\n", "\n", - "crossmatched_objects = filter_functions.catalog_crossmatch(\n", + "crossmatched_objects = catalog_crossmatch(\n", " alerts=filtered_objects,\n", - " catalog_name=[\n", - " \"CLAUDS_SourceExtractor_XMMLSS_cut.fits\",\n", - " \"CLAUDS_SourceExtractor_DEEP23_cut.fits\"\n", - " ],\n", ")" ] }, @@ -337,7 +275,7 @@ "source": [ "# additional filtering on catalog features\n", "\n", - "from filter_functions import catalog_filter\n", + "from thor.utils.filter_functions import catalog_filter\n", "\n", "filtered = catalog_filter(\n", " crossmatched_objects,\n", @@ -353,11 +291,13 @@ "outputs": [], "source": [ "# Source specific filtering:\n", - "# in tde_filter we also make more specific cuts, including minimum x detections (specfic to DDFs), no milliquas matches\n", + "# in tde_filter we also make more specific cuts, including minimum # detections, no milliquas matches\n", "\n", - "filtered = filter_functions.filter_alerts( \n", + "from thor.utils.filter_functions import (filter_alerts, tde_filter)\n", + "\n", + "filtered = filter_alerts( \n", " filtered, \n", - " filter_functions.tde_filter,\n", + " tde_filter\n", ") " ] }, @@ -377,10 +317,10 @@ "outputs": [], "source": [ "crossmatched_alerts = [v[\"LSST\"] for v in crossmatched_objects.values()]\n", - "print(f\"Total alerts: {len(crossmatched_alerts)}\")\n", + "print(f\"Total crossmatched alerts: {len(crossmatched_alerts)}\")\n", "\n", "filtered_crossmatched_alerts = [v[\"LSST\"] for v in filtered.values()] \n", - "print(f\"Filtered alerts: {len(filtered_crossmatched_alerts)}\")" + "print(f\"Filtered crossmatched alerts: {len(filtered_crossmatched_alerts)}\")" ] }, { @@ -392,9 +332,10 @@ "source": [ "# inspect all crossmatch results without the cuts on rise etc.\n", "\n", + "from babamul.jupyter import scan_alerts\n", "from astropy.time import Time\n", "print(f'Todays JD: {round(Time.now().jd)}')\n", - "babamul.jupyter.scan_alerts(filtered_crossmatched_alerts)" + "scan_alerts(filtered_crossmatched_alerts)" ] }, { @@ -404,7 +345,7 @@ "metadata": {}, "outputs": [], "source": [ - "# print crossmatch info for a given id \n", + "# print crossmatch info for a given id from dictionary\n", "\n", "id = '170591539645907039'\n", "\n", @@ -414,7 +355,8 @@ " print(f\"{catalog}:\")\n", " print(f\"end=z = {data['z']:.2f}\") \n", " print(f\"chi2 = {data['CHI_BEST']:.2f}\")\n", - " print(f\"separation = {data['consearch_arcsecs']:.2f} arcsec\")\n" + " if \"consearch_arcsecs\" in data:\n", + " print(f\"separation = {data['consearch_arcsecs']:.2f} arcsec\")\n" ] }, { @@ -429,8 +371,9 @@ "\n", "cand_ids_to_save = []\n", "\n", - "from rubin_stats_functions import save_objects \n", - "save_objects(cand_ids_to_save, path=\"data/tde_candidates.json.gz\")" + "from thor.utils.fetch_alerts import save_objects\n", + "\n", + "save_objects(cand_ids_to_save, path=\"../../data/lsst_alert_download/tde_candidates.json.gz\")" ] }, { @@ -451,11 +394,13 @@ "# or we provide function to run cosmos crossmatch by directly provided ra and dec instead of alert\n", "# you may provide a single ra and dec, or a list of ras and a list of corresponding decs as arguments.\n", "\n", - "crossmatched_coords = filter_functions.catalog_crossmatch(\n", + "from thor.utils.filter_functions import catalog_crossmatch\n", + "\n", + "crossmatched_coords = catalog_crossmatch(\n", " ra=150.101273,\n", " dec=2.743894,\n", " radius_arcsec=300,\n", - " catalog_path=\"data/catalogs\",\n", + " catalog_path=\"../../data/catalogs\",\n", ")" ] }, @@ -468,11 +413,29 @@ "source": [ "# or we can load our list of previously saved candidates, and fetch alerts to display\n", "\n", - "from rubin_stats_functions import load_objects, fetch_latest_alerts \n", + "from thor.utils.filter_functions import catalog_crossmatch\n", + "from thor.utils.fetch_alerts import (load_objects, fetch_latest_alerts)\n", + "\n", + "object_ids = load_objects(\"../../data/lsst_alert_download/tde_candidates.json.gz\")\n", + "latest_alerts = fetch_latest_alerts(object_ids, save=False, verbose=False)\n", + "saved_objects = catalog_crossmatch(latest_alerts)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d65f7e61", + "metadata": {}, + "outputs": [], + "source": [ + "# look at our candidates:\n", + "\n", + "from babamul.jupyter import scan_alerts\n", + "from astropy.time import Time\n", + "print(f'Todays JD: {round(Time.now().jd)}')\n", "\n", - "object_ids = load_objects(\"data/tde_candidates.json.gz\") \n", - "latest_alerts = fetch_latest_alerts(object_ids) \n", - "saved_objects = filter_functions.catalog_crossmatch(latest_alerts) " + "show_saved_objects = [v[\"LSST\"] for v in saved_objects.values()] \n", + "scan_alerts(show_saved_objects)" ] }, { @@ -480,7 +443,7 @@ "id": "9561b1bc", "metadata": {}, "source": [ - "### prost\n", + "### prost - currently developing bayesian crossmatch instead of basic conesearch\n", "\n", "https://github.com/alexandergagliano/Prost\n", "\n", @@ -493,6 +456,18 @@ "can return top 2 best host associations." ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "834d5a18", + "metadata": {}, + "outputs": [], + "source": [ + "import importlib\n", + "import thor.utils.fetch_alerts\n", + "importlib.reload(thor.utils.fetch_alerts)" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/docs/notebooks/Example_Visualizing_LSST_Observations.ipynb b/docs/notebooks/Example_Visualizing_LSST_Observations.ipynb index 35a3dd7..f969e8b 100644 --- a/docs/notebooks/Example_Visualizing_LSST_Observations.ipynb +++ b/docs/notebooks/Example_Visualizing_LSST_Observations.ipynb @@ -5,7 +5,7 @@ "id": "14a81a74", "metadata": {}, "source": [ - "# Get statistics and visualizations of LSST alert stream and Rubin Observations\n", + "# Get statistics and visualizations of LSST alert stream and Rubin observations\n", "\n", "### Other useful resources include:\n", "\n", @@ -41,18 +41,6 @@ "dotenv.load_dotenv()" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "961e78f1", - "metadata": {}, - "outputs": [], - "source": [ - "import importlib\n", - "import thor.utils.rubin_stats_visualizations\n", - "importlib.reload(thor.utils.rubin_stats_visualizations)" - ] - }, { "cell_type": "code", "execution_count": null, diff --git a/src/thor/crossmatch_alerts.py b/src/thor/crossmatch_alerts.py index b49b81e..abeb2c2 100644 --- a/src/thor/crossmatch_alerts.py +++ b/src/thor/crossmatch_alerts.py @@ -1,70 +1,207 @@ import argparse import json import os -import sys +import subprocess +import tempfile from datetime import datetime, timezone from pathlib import Path -from thor.utils.filter_functions import filter_alerts, generic_filter, tde_filter -from thor.utils.fetch_alerts import load_alerts +import dotenv + +from thor.utils import filter_functions +from thor.utils.fetch_alerts import babamul_get_alerts + +Z_COLS = {"z", "Z_BEST", "ZPHOT", "zfinal", "zpdf_med"} + + +def _launch_scan_notebook(object_ids): + """Write a temporary notebook pre-loaded with candidates and open it.""" + object_ids_repr = json.dumps(list(object_ids)) + nb = { + "nbformat": 4, + "nbformat_minor": 5, + "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, + "cells": [ + { + "cell_type": "code", + "id": "scan-setup", + "metadata": {}, + "outputs": [], + "source": ( + "import dotenv\n" + "dotenv.load_dotenv()\n" + "\n" + "import babamul.api as _api\n" + "from babamul import LsstAlert\n" + "from babamul.jupyter import scan_alerts\n" + "\n" + f"object_ids = {object_ids_repr}\n" + "\n" + "alerts = []\n" + "for oid in object_ids:\n" + " try:\n" + " alerts.append(_api.get_object('LSST', oid))\n" + " except Exception as e:\n" + " print(f'Failed {oid}: {e}')\n" + "\n" + "print(f'Loaded {len(alerts)} alerts.')\n" + "scan_alerts(alerts)" + ), + } + ], + } + + tmp = tempfile.NamedTemporaryFile( + suffix=".ipynb", prefix="thor_scan_", delete=False, mode="w" + ) + json.dump(nb, tmp) + tmp.close() + print(f"\nOpening scan notebook — press Ctrl+C in this terminal when done to close and clean up.") + try: + subprocess.run(["jupyter", "notebook", tmp.name], env=os.environ.copy()) + finally: + os.unlink(tmp.name) + + +def _print_match_report(crossmatched_objects): + n = len(crossmatched_objects) + print(f"\nMatched Object IDs: {n}") + + if n > 100: + print(f"More than 100 matches ({n} total) — skipping per-object summary.") + return + + # column widths + id_w = max(len("LSST Object ID"), max(len(str(oid)) for oid in crossmatched_objects)) + cat_w = 30 + z_w = 8 + sep_w = 10 + + header = ( + f"{'LSST Object ID':<{id_w}} " + f"{'Catalog':<{cat_w}} " + f"{'z':>{z_w}} " + f"{'Sep (\")':>{sep_w}}" + ) + divider = "-" * len(header) + print(divider) + print(header) + print(divider) + + for obj_id, obj in crossmatched_objects.items(): + first = True + for catalog, data in obj.items(): + if catalog == "LSST" or data is None: + continue + z_val = next((data[c] for c in Z_COLS if c in data and data[c] is not None), None) + sep_val = data.get("conesearch_arcsecs") + z_str = f"{z_val:.3f}" if z_val is not None else "—" + sep_str = f"{sep_val:.2f}" if sep_val is not None else "—" + id_str = str(obj_id) if first else "" + print( + f"{id_str:<{id_w}} " + f"{catalog:<{cat_w}} " + f"{z_str:>{z_w}} " + f"{sep_str:>{sep_w}}" + ) + first = False + + print(divider) def main(): - parser = argparse.ArgumentParser(description="Run TDE filter pipeline on local alert data.") + dotenv.load_dotenv() + parser = argparse.ArgumentParser(description="Fetch LSST alerts and crossmatch against catalogs.") + parser.add_argument("--start", required=True, help="Start date (MM-DD-YYYY).") + parser.add_argument("--end", required=True, help="End date (MM-DD-YYYY).") + parser.add_argument( + "--additional_filtering", + default=None, + choices=["tde_filter"], + help="Optional additional filter to apply after crossmatch (default: none).", + ) parser.add_argument( - "--test_data", - required=True, - help="Path to directory containing .json.gz alert files.", + "--save", + action="store_true", + help="Save fetched alerts to data/lsst_alert_download/raw_files/ (default: off).", + ) + parser.add_argument( + "--save_result", + action="store_true", + help="Save crossmatch results to data/lsst_alert_download/ (default: off).", + ) + parser.add_argument( + "--scan", + action="store_true", + help="Open a temporary Jupyter notebook to scan candidates with scan_alerts (default: off).", ) args = parser.parse_args() - data_dir = Path(args.test_data) - if not data_dir.is_dir(): - raise ValueError(f"--test_data must be a directory, got: {data_dir}") - - # ── Load alerts from all .json.gz files in the directory ───────────────── - alert_files = sorted(data_dir.glob("*.json.gz")) - if not alert_files: - raise FileNotFoundError(f"No .json.gz files found in {data_dir}") - - loaded_alerts = [] - for f in alert_files: - loaded_alerts.extend(load_alerts(f, survey="LSST")) - print(f"Loaded {len(loaded_alerts):,} alerts from {len(alert_files)} file(s).") + # ── Fetch alerts ────────────────────────────────────────────────────────── + loaded_alerts = babamul_get_alerts( + survey="LSST", + start_time=args.start, + end_time=args.end, + min_drb=0.4, + is_rock=False, + is_star=False, + is_near_brightstar=False, + is_stationary=True, + save=args.save, + ) + print(f"Loaded {len(loaded_alerts):,} alerts.") - # ── Generic cuts (drb, rock, star, near_brightstar, stationary, PSF) ───── - filtered_alerts = filter_alerts( + # ── Generic cuts ────────────────────────────────────────────────────────── + filtered_alerts = filter_functions.filter_alerts( loaded_alerts, - generic_filter, + filter_functions.generic_filter, ) - # ── TDE-specific cuts (>=5 detections, no Milliquas match, rising) ──────── - tde_candidates = filter_alerts( - filtered_alerts, - tde_filter, + # ── Deduplicate to unique objects ───────────────────────────────────────── + filtered_objects = filter_functions.deduplicate_alerts(filtered_alerts) + + # ── Crossmatch against all available catalogs ───────────────────────────── + crossmatched_objects = filter_functions.catalog_crossmatch( + alerts=filtered_objects, ) - # ── Save results ────────────────────────────────────────────────────────── + # ── Optional additional filtering ───────────────────────────────────────── + if args.additional_filtering == "tde_filter": + crossmatched_objects = filter_functions.filter_alerts( + crossmatched_objects, + filter_functions.tde_filter, + ) + + # ── Report ──────────────────────────────────────────────────────────────── + if not crossmatched_objects: + print("\nNo crossmatch candidates found.") + return + + _print_match_report(crossmatched_objects) + + # ── Optionally scan in Jupyter ──────────────────────────────────────────── + if args.scan: + _launch_scan_notebook(list(crossmatched_objects.keys())) + + # ── Optionally save results ─────────────────────────────────────────────── + if not args.save_result: + return + + serializable = { + obj_id: {k: v for k, v in obj.items() if k != "LSST"} + for obj_id, obj in crossmatched_objects.items() + } + timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%dT%H%M%SZ") - out_dir = Path("filter_results") / f"filter_test_{timestamp}" + repo_root = Path(__file__).resolve().parents[2] + out_dir = repo_root / "data" / "lsst_alert_download" out_dir.mkdir(parents=True, exist_ok=True) + out_file = out_dir / f"crossmatch_candidates_{timestamp}.json" - results = { - "timestamp": timestamp, - "test_data": str(data_dir), - "n_loaded": len(loaded_alerts), - "n_after_generic": len(filtered_alerts), - "n_tde_candidates": len(tde_candidates), - "candidates": [a.objectId for a in tde_candidates], - } - out_file = out_dir / "results.json" with open(out_file, "w") as f: - json.dump(results, f, indent=2) + json.dump(serializable, f, indent=2, default=str) - print(f"\nResults saved to {out_file}") - print(f" Loaded: {results['n_loaded']:,}") - print(f" After generic: {results['n_after_generic']:,}") - print(f" TDE candidates: {results['n_tde_candidates']:,}") + print(f"\nSaved {len(serializable):,} crossmatch candidates to {out_file}") if __name__ == "__main__": diff --git a/src/thor/utils/fetch_alerts.py b/src/thor/utils/fetch_alerts.py index c99f0b1..70a921b 100644 --- a/src/thor/utils/fetch_alerts.py +++ b/src/thor/utils/fetch_alerts.py @@ -11,13 +11,16 @@ from pydantic import ValidationError _LIST_FIELDS = frozenset(("prv_candidates", "fp_hists")) -_FLOAT_SUFFIXES = ("chi2", "rate", "rate_error", "dt", "jd", "mag", "err") +_FLOAT_SUFFIXES = ("chi2", "rate", "rate_error", "dt", "jd", "mag", "err", "Err", "flux", "Flux") +# Defaults injected for required fields added to the babamul schema after old alerts were saved. +_MISSING_BOOL_DEFAULTS = {"isNegative": False} def _coerce_alert(obj): """Recursively fix API nulls that babamul models don't allow: - None list fields → [] - None float fields (by name heuristic) → nan + - Missing required bool fields (schema additions) → default value """ if isinstance(obj, list): return [_coerce_alert(i) for i in obj] @@ -31,9 +34,40 @@ def _coerce_alert(obj): out[k] = float("nan") else: out[k] = _coerce_alert(v) + for field, default in _MISSING_BOOL_DEFAULTS.items(): + if field not in out: + out[field] = default return out +# ── Monkey-patch babamul.api.get_object to coerce nulls before model validation ── +# babamul's lazy photometry fetch calls get_object internally. Some LSST alerts +# return null for fields like psfFluxErr that the model requires to be floats. +# Wrapping get_object with _coerce_alert fixes these the same way we fix saved alerts. +import babamul.api as _babamul_api + +_orig_get_object = _babamul_api.get_object + +def _patched_get_object(survey, object_id): + import base64 as _base64 + from babamul import LsstAlert as _LsstAlert, ZtfAlert as _ZtfAlert + from babamul.api import _request, get_args, Survey + response = _request("GET", f"/surveys/{survey}/objects/{object_id}") + data = response.get("data", response) + for key in ["cutoutScience", "cutoutTemplate", "cutoutDifference"]: + if data.get(key) and isinstance(data[key], str): + data[key] = _base64.b64decode(data[key]) + data = _coerce_alert(data) + if survey == "ZTF": + return _ZtfAlert.model_validate(data) + elif survey == "LSST": + return _LsstAlert.model_validate(data) + else: + valid_surveys = ", ".join(get_args(Survey)) + raise ValueError(f"Survey {survey} is not supported, must be one of: {valid_surveys}") + +_babamul_api.get_object = _patched_get_object + # ── Alert fetching ───────────────────────────────────────────────────────────── def babamul_get_alerts( @@ -330,32 +364,44 @@ def load_objects(path): return ids -def fetch_latest_alerts(object_ids, survey="LSST"): +def fetch_latest_alerts(object_ids, survey="LSST", save=True, verbose=True): """ - Fetch the most recent alert for each object ID. + Fetch the current object record for each object ID using a single API call + per object (get_object), which is much faster than pulling full alert history. Parameters ---------- object_ids : list of str survey : str "LSST" or "ZTF". Default "LSST". + save : bool + If True (default), save the fetched alerts to + data/lsst_alert_download/raw_files/latest_alerts_{survey}.json.gz. + verbose : bool + If True (default), print progress per object. Returns ------- list of ZtfAlert | LsstAlert - One alert per object ID (the most recent by JD/MJD). + One alert per successfully fetched object ID. """ latest = [] + failed = [] for oid in object_ids: - alerts = babamul_get_alerts(survey=survey, object_id=oid) - if not alerts: - continue - most_recent = max( - alerts, - key=lambda a: a.candidate.jd if hasattr(a.candidate, "jd") else a.candidate.midpointMjdTai, - ) - latest.append(most_recent) - print(f"Fetched latest alert for {len(latest):,}/{len(object_ids):,} objects.") + if verbose: + print(f"Fetching {oid}...") + try: + alert = _babamul_api.get_object(survey, oid) + latest.append(alert) + except Exception as e: + failed.append(oid) + if verbose: + print(f" Failed for {oid}: {e}") + print(f"Fetched {len(latest):,}/{len(object_ids):,} objects ({len(failed):,} failed).") + if save and latest: + _repo_root = Path(__file__).resolve().parents[3] + out_path = _repo_root / "data" / "lsst_alert_download" / "raw_files" / f"latest_alerts_{survey}.json.gz" + save_alerts(latest, out_path) return latest @@ -383,7 +429,16 @@ def _object_hook(obj): data = json.load(f, object_hook=_object_hook) print(f"Loaded {len(data):,} alerts from {path}.") - return [model.model_validate(_coerce_alert(a)) for a in data] + alerts = [model.model_validate(_coerce_alert(a)) for a in data] + # The main alerts API doesn't return photometry history inline — those fields + # are fetched lazily by babamul when show() is called. _coerce_alert converts + # the None values to [] so pydantic accepts them, but that tricks babamul into + # thinking photometry was already fetched. Reset to None to re-enable lazy fetch. + for alert in alerts: + if not alert.prv_candidates and not alert.fp_hists: + alert.prv_candidates = None + alert.fp_hists = None + return alerts def combine_alert_files(input_dir, output_path, pattern="*.json.gz", input_files=None, delete_raw=False):