From 07daa7e42548f1a76e967091ac57d6327fe4b86d Mon Sep 17 00:00:00 2001 From: Kira Date: Thu, 2 Jul 2026 08:52:45 -0700 Subject: [PATCH] organize redback work --- .gitignore | 2 + .../Develop_Redback_Simulations.ipynb | 233 +++++++----------- src/thor/crossmatch_alerts_fritz.py | 0 src/thor/utils/redback_sims.py | 232 +++++++++++++++++ 4 files changed, 317 insertions(+), 150 deletions(-) delete mode 100644 src/thor/crossmatch_alerts_fritz.py diff --git a/.gitignore b/.gitignore index f7c9f0f..e8a3472 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ # data (catalogs, simulations, rubin alerts) data/ +docs/notebooks/tidal_disruption_event/ +docs/notebooks/redback_inference # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/docs/notebooks/Develop_Redback_Simulations.ipynb b/docs/notebooks/Develop_Redback_Simulations.ipynb index bd5d6d0..9e5e90f 100644 --- a/docs/notebooks/Develop_Redback_Simulations.ipynb +++ b/docs/notebooks/Develop_Redback_Simulations.ipynb @@ -4,11 +4,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# redback simulations\n", + "# Redback Simulations\n", "\n", - "use redback to simulate tdes and other transients at different redshifts\n", - "\n", - "https://redback.readthedocs.io/en/latest/simulation.html#examples" + "Use [redback](https://redback.readthedocs.io/en/latest/simulation.html#examples) to simulate TDEs and other transients at different redshifts." ] }, { @@ -17,6 +15,13 @@ "metadata": {}, "outputs": [], "source": [ + "import warnings\n", + "# Suppress SyntaxWarnings from invalid escape sequences in redback/bilby source files\n", + "warnings.filterwarnings('ignore', category=SyntaxWarning)\n", + "# Suppress optional-dependency warnings from redback (lalsimulation, swifttools)\n", + "warnings.filterwarnings('ignore', message='.*lalsimulation.*')\n", + "warnings.filterwarnings('ignore', message='.*swifttools.*')\n", + "\n", "import numpy as np\n", "import pandas as pd\n", "import os\n", @@ -27,9 +32,7 @@ "import matplotlib \n", "from matplotlib import rcParams\n", "rcParams[\"font.family\"] = \"Liberation Serif\"\n", - "matplotlib.rcParams['text.usetex'] = False \n", - "\n", - "# FIXME LALsimulation not installed" + "matplotlib.rcParams['text.usetex'] = False" ] }, { @@ -38,9 +41,7 @@ "source": [ "## Fetch real object photometry\n", "\n", - "https://redback.readthedocs.io/en/latest/getting_data.html\n", - "\n", - "redback has built in tools to get data from open access, otter catalogs and fink and lasair brokers" + "redback has [built in tools to get data](https://redback.readthedocs.io/en/latest/getting_data.html) from open access, otter catalogs and fink / lasair brokers" ] }, { @@ -73,7 +74,7 @@ "tdes['maxabsmag'] = pd.to_numeric(tdes['maxabsmag'], errors='coerce') \n", " \n", "# print brightest \n", - "brightest = tdes[['name', 'maxabsmag']].dropna().sort_values('maxabsmag').head(20) \n", + "brightest = tdes[['name', 'maxabsmag']].dropna().sort_values('maxabsmag').head(10) \n", "print(\"Brightest TDEs by peak absolute magnitude:\") \n", "print(brightest.to_string(index=False)) \n", " \n", @@ -94,7 +95,7 @@ "metadata": {}, "outputs": [], "source": [ - "# instead of filtering for OAC TDEs, just check brightest objects\n", + "# instead of filtering for OAC TDEs, just check brightest objects in general\n", "\n", "oac_path = os.path.join(os.path.dirname(redback.__file__), 'tables', 'OAC_metadata.csv') \n", "df = pd.read_csv(oac_path, low_memory=False, on_bad_lines='skip')\n", @@ -109,10 +110,12 @@ "metadata": {}, "outputs": [], "source": [ + "# we can also use Otter to query for objects\n", + "\n", "import otter\n", "import numpy as np\n", "\n", - "db = otter.Otter()\n", + "db = otter.Otter(datadir='../../data/redback_sims')\n", "\n", "# query for confirmed TDEs with optical photometry in one step\n", "all_transients = db.query(classification='TDE', unambiguous=True,)# has_uvoir_phot=True)\n", @@ -132,7 +135,7 @@ "\n", "tdes_with_z.sort(key=lambda x: x[1], reverse=True)\n", "\n", - "print(f\"Found {len(tdes_with_z)} confirmed TDEs with UVOIR photometry\\n\")\n", + "print(f\"Found {len(tdes_with_z)} confirmed TDEs\\n\")\n", "print(f\"{'Rank':<6} {'Name':<25} {'Redshift'}\")\n", "print(\"-\" * 45)\n", "for i, (name, z) in enumerate(tdes_with_z[:15], 1):\n", @@ -145,7 +148,9 @@ "metadata": {}, "outputs": [], "source": [ - "target_object = '2025aarm'\n", + "# check for a specific transient in the results by name\n", + "\n", + "target_object = '2019cmw'\n", "\n", "for t in all_transients:\n", " name = t.get('name', {}).get('default_name', str(t))\n", @@ -158,7 +163,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### now load photometry for selected real objects" + "### now load photometry into redback structure for selected real objects" ] }, { @@ -167,27 +172,9 @@ "metadata": {}, "outputs": [], "source": [ - "# FIXME redback requires to make a directory\n", - "# also can try AT2018iih (z=0.212), AT2020ysg, AT2020acka\n", - "\n", - "# load with name\n", - "name = 'AT2019cmw'\n", - "\n", - "# redback bug fix\n", - "outdir = f'tidal_disruption_event/{name}/uvoir'\n", - "os.makedirs(outdir, exist_ok=True)\n", - "\n", - "# either load from otter (tdes only right now)\n", - "data = redback.get_data.get_tidal_disruption_event_data_from_otter(transient=name)\n", - "\n", - "# # or load tdes from open access catalog\n", - "# data = redback.get_data.get_tidal_disruption_event_data_from_open_transient_catalog_data(transient=name)\n", - "\n", - "# # or load SN from OAC\n", - "# data = redback.get_data.get_supernova_data_from_open_transient_catalog_data(transient=name)\n", - "\n", - "# # or query Fink for ZTF objects\n", - "# data = redback.get_data.get_fink_data(transient=name, transient_type='supernova')" + "import importlib\n", + "import thor.utils.redback_sims\n", + "importlib.reload(thor.utils.redback_sims)" ] }, { @@ -196,12 +183,12 @@ "metadata": {}, "outputs": [], "source": [ - "# can query FINK and LASAIR also\n", - "#FIXME fetch more bands of data.\n", - "\n", - "name = 'ZTF25acfyeke'\n", + "# Load with name, for example AT2019cmw, AT2018iih (z=0.212), AT2020ysg, AT2020acka\n", + "name = 'AT2020qhs'\n", + "source = 'otter'\n", "\n", - "data = redback.get_data.get_fink_data(transient=name, transient_type='tidal_disruption_event')" + "from thor.utils.redback_sims import load_data_redback_object, create_redback_transient\n", + "data = load_data_redback_object(object_name=name, data_source=source)" ] }, { @@ -211,21 +198,16 @@ "outputs": [], "source": [ "# create redback transient object from loaded data\n", - "exclude_bands = ['W2', 'M2', 'W1', 'U']\n", - "bands=data['band'].unique() # utilize all bands available\n", - "#bands = [b for b in bands if b not in exclude_bands] # filter out non-optical bands\n", - "\n", - "# # from OTTER TDE (also downloads if not cached) \n", - "# transient_object = redback.tde.TDE.from_otter(name=name, data_mode=\"flux_density\") \n", - "\n", - "# # from OAC TDE \n", - "# transient_object = redback.tde.TDE.from_open_access_catalogue(name=name, data_mode=\"flux_density\", active_bands=bands)\n", - "\n", - "# # from OAC Supernova \n", - "# transient_object = redback.supernova.Supernova.from_open_access_catalogue(name=name, data_mode=\"flux_density\", active_bands=bands) \n", - " \n", - "# from Fink — same as OAC since they share the same processed file path \n", - "transient_object = redback.tde.TDE.from_open_access_catalogue(name=name, data_mode=\"flux_density\", active_bands=bands) " + "exclude_bands = ['W2', 'M2', 'W1', 'U'] # set exclude_bands = None # utilize all bands available\n", + "\n", + "transient_object = create_redback_transient(\n", + " object_name=name,\n", + " data_source=source,\n", + " data=data,\n", + " data_mode='flux_density',\n", + " exclude_bands=exclude_bands,\n", + " plot=True\n", + ")" ] }, { @@ -236,17 +218,17 @@ "source": [ "# plot\n", "\n", - "# the black points are bands available that havent be set as active \n", - "transient_object.plot_data(show=False, title=f\"{name} photometry\") \n", - "# fig, axes = plt.subplots(1,3, sharex=True, sharey=True, figsize=(12, 5)) \n", - "# object.plot_multiband(figure=fig, axes=axes, filters=[\"g'\", \"r'\", \"z'\"], show=False) " + "# the black points are bands available that havent been set as active\n", + "transient_object.plot_data(show=False, title=f\"{name} photometry\")\n", + "fig, axes = plt.subplots(1,3, sharex=True, sharey=True, figsize=(12, 5))\n", + "transient_object.plot_multiband(figure=fig, axes=axes, filters=[\"g\", \"r\", \"UVW1\"], show=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Warning: \n", + "Usage warning: \n", "\n", "- check if data are in AB or Vega system (redback assumes AB, throws away explicity flagged Vega from processed file)\n", "\n", @@ -256,31 +238,26 @@ ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "# FIXME automate data checks\n", - "\n", - "# check what bands the observations are in\n", - "print(\"Unique bands:\", transient_object.unique_bands) \n", - "print(\"Active bands:\", transient_object.active_bands) \n", - "\n", - "# count photometry points per band\n", - "from collections import Counter\n", - "all_counts = Counter(transient_object.bands)\n", - "print(\"\\nAll bands (including inactive):\")\n", - "for band, count in sorted(all_counts.items()):\n", - " active = band in transient_object.active_bands\n", - " print(f\" {band}: {count}\")" + "# fit model based on real photometry" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "# fit model based on real photometry" + "### [Redback Models](https://redback.readthedocs.io/en/latest/models.html)\n", + "\n", + "tde model options:\n", + "Simple analytic fallback\n", + "Cooling envelope (Metzger 2022)\n", + "Gaussian rise + cooling envelope (Sarin and Metzger 2024)\n", + "Broken power law + cooling envelope (Sarin and Metzger 2024)\n", + "Fitted\n", + "TDEMass\n", + "Mosfit TDE" ] }, { @@ -289,16 +266,10 @@ "metadata": {}, "outputs": [], "source": [ - "# run inference to extract model best parameters\n", + "# print all the priors for a given model\n", "\n", - "# tde model options:\n", - "# Simple analytic fallback\n", - "# Cooling envelope (Metzger 2022)\n", - "# Gaussian rise + cooling envelope (Sarin and Metzger 2024)\n", - "# Broken power law + cooling envelope (Sarin and Metzger 2024)\n", - "# Fitted\n", - "# TDEMass\n", - "# Mosfit TDE" + "priors = redback.priors.get_priors(model='tde_analytical')\n", + "priors" ] }, { @@ -307,73 +278,35 @@ "metadata": {}, "outputs": [], "source": [ - "sampler = 'dynesty'\n", - "model = 'tde_analytical'\n", - "nlive=500\n", - "outdir = f'redback_inference/{name}/{model}'\n", - "os.makedirs(outdir, exist_ok=True)\n", - "\n", - "# # By default all data is used, select bands for speed or if you only trust some of the data/physics.\n", - "# bands = [\"u\", \"r\"]\n", - "# transient_object.active_bands = bands\n", + "# Fit model for known object. \n", + "# depending on priors etc, this can take 20+ minutes to run.\n", "\n", - "# start with default priors\n", - "priors = redback.priors.get_priors(model=model)\n", - "# Fix priors, ie if known through the metadata that was downloaded alongside the data\n", - "priors['redshift'] = 0.4030\n", - "# tighten temperature prior — TDE blackbodies are typically 10k-50k K at peak\n", - "priors['temperature_0'] = bilby.core.prior.Uniform(\n", - " minimum=10000, maximum=60000, name='temperature_0', latex_label='$T_0$ (K)')\n", - "# tighten radius prior around physically expected TDE photosphere size\n", - "priors['radius_0'] = bilby.core.prior.LogUniform(\n", - " minimum=1e13, maximum=1e16, name='radius_0', latex_label='$R_0$ (cm)')\n", - "\n", - "model_kwargs = dict(frequency=transient_object.filtered_frequencies, output_format='flux_density')\n", - "\n", - "# returns a tde result object\n", - "result = redback.fit_model(transient=transient_object, model=model, sampler=sampler,\n", - " model_kwargs=model_kwargs, prior=priors, sample='rslice', nlive=nlive, npool=4, \n", - " resume=True, outdir=outdir, label=f\"{name}_nlive{str(nlive)}\",\n", - " plot=False)\n", - "\n", - "# result.transient reconstructs from metadata with wrong path, use transient_object directly\n", - "model_func = redback.model_library.all_models_dict[model]\n", + "import bilby\n", + "import importlib\n", + "import thor.utils.redback_sims\n", + "importlib.reload(thor.utils.redback_sims)\n", + "from thor.utils.redback_sims import fit_redback_model\n", + "\n", + "prior_overrides = {\n", + " 'redshift': 0.4030,\n", + " # tighten temperature prior — TDE blackbodies are typically 10k-50k K at peak\n", + " 'temperature_0': bilby.core.prior.Uniform(\n", + " minimum=10000, maximum=60000, name='temperature_0', latex_label='$T_0$ (K)'),\n", + " # tighten radius prior around physically expected TDE photosphere size\n", + " 'radius_0': bilby.core.prior.LogUniform(\n", + " minimum=1e13, maximum=1e16, name='radius_0', latex_label='$R_0$ (cm)'),\n", + "}\n", "\n", - "# plot results\n", - "result.plot_corner(save=False, show=True)\n", - "plt.show()\n", - "plt.close('all')\n", - "transient_object.plot_lightcurve(\n", - " model=model_func,\n", - " posterior=result.posterior,\n", - " model_kwargs=result.model_kwargs,\n", - " random_models=100,\n", - " save=False,\n", - " show=True,\n", - ")\n", - "transient_object.plot_multiband_lightcurve(\n", - " model=model_func,\n", - " posterior=result.posterior,\n", - " model_kwargs=result.model_kwargs,\n", - " random_models=100,\n", - " filters=[\"ztfg\", \"ztfr\"],\n", - " save=False,\n", - " show=True,\n", + "result = fit_redback_model(\n", + " transient_object=transient_object,\n", + " object_name=name,\n", + " model='tde_analytical',\n", + " nlive=500,\n", + " prior_overrides=prior_overrides,\n", + " multiband_filters=['ztfg', 'ztfr'],\n", ")" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# print all the priors for the fit model\n", - "\n", - "priors = redback.priors.get_priors(model='tde_analytical')\n", - "priors" - ] - }, { "cell_type": "code", "execution_count": null, diff --git a/src/thor/crossmatch_alerts_fritz.py b/src/thor/crossmatch_alerts_fritz.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/thor/utils/redback_sims.py b/src/thor/utils/redback_sims.py index e69de29..9d605fe 100644 --- a/src/thor/utils/redback_sims.py +++ b/src/thor/utils/redback_sims.py @@ -0,0 +1,232 @@ +import redback +import os +import matplotlib.pyplot as plt + +VALID_SOURCES = ('otter', 'oac_tde', 'oac_sn', 'fink_tde', 'fink_sn') + + +def load_data_redback_object(object_name, data_source): + """ + Load photometry for a transient via redback and return the data table. + + Parameters + ---------- + object_name : str + Transient name (e.g. 'AT2019cmw', 'ZTF25acfyeke'). + data_source : str + One of: + 'otter' – TDE from the OTTER catalog + 'oac_tde' – TDE from the Open Access Catalog + 'oac_sn' – supernova from the Open Access Catalog + 'fink_tde' – TDE from Fink (ZTF stream) + 'fink_sn' – supernova from Fink (ZTF stream) + + Returns + ------- + data : pd.DataFrame + Photometry table returned by redback. + """ + if data_source not in VALID_SOURCES: + raise ValueError(f"data_source must be one of {VALID_SOURCES}, got {data_source!r}") + + # redback bug fix: otter requires this directory to exist before fetching + if data_source == 'otter': + uvoir_dir = f'tidal_disruption_event/{object_name}/uvoir' + os.makedirs(uvoir_dir, exist_ok=True) + # if a previous run failed after writing the raw file but before writing + # metadata, delete the raw file so both get re-fetched cleanly + raw = f'{uvoir_dir}/{object_name}_rawdata.csv' + meta = f'{uvoir_dir}/{object_name}_metadata.csv' + if os.path.exists(raw) and not os.path.exists(meta): + os.remove(raw) + + if data_source == 'otter': + return redback.get_data.get_tidal_disruption_event_data_from_otter(transient=object_name) + elif data_source == 'oac_tde': + return redback.get_data.get_tidal_disruption_event_data_from_open_transient_catalog_data(transient=object_name) + elif data_source == 'oac_sn': + return redback.get_data.get_supernova_data_from_open_transient_catalog_data(transient=object_name) + elif data_source == 'fink_tde': + return redback.get_data.get_fink_data(transient=object_name, transient_type='tidal_disruption_event') + elif data_source == 'fink_sn': + return redback.get_data.get_fink_data(transient=object_name, transient_type='supernova') + + +def create_redback_transient(object_name, data_source, data, + data_mode='flux_density', exclude_bands=None): + """ + Create a redback transient object from previously loaded photometry data. + + Parameters + ---------- + object_name : str + Transient name (e.g. 'AT2019cmw'). + data_source : str + Same source used in load_data_redback_object. Controls which redback + class and constructor are used. + data : pd.DataFrame + Output of load_data_redback_object — used to determine available bands. + data_mode : str + Redback data mode. Default: 'flux_density'. + exclude_bands : list or None + Bands to exclude from active_bands (e.g. ['W2', 'M2', 'W1', 'U']). + If None, all bands in the data are used. + + Returns + ------- + transient_object : redback transient + Redback TDE or Supernova object ready for fitting/plotting. + """ + if data_source not in VALID_SOURCES: + raise ValueError(f"data_source must be one of {VALID_SOURCES}, got {data_source!r}") + + bands = data['band'].unique() + if exclude_bands is not None: + bands = [b for b in bands if b not in exclude_bands] + + if data_source == 'otter': + return redback.tde.TDE.from_otter( + name=object_name, data_mode=data_mode, active_bands=bands) + elif data_source in ('oac_tde', 'fink_tde'): + return redback.tde.TDE.from_open_access_catalogue( + name=object_name, data_mode=data_mode, active_bands=bands) + elif data_source in ('oac_sn', 'fink_sn'): + return redback.supernova.Supernova.from_open_access_catalogue( + name=object_name, data_mode=data_mode, active_bands=bands) + + +def _sort_transient_by_time(transient_object): + """Sort all per-observation arrays on a redback transient in-place by time. + + Workaround for a redback bug in Diffusion.convert_input_luminosity where + dense_times is built from time[-1] rather than time.max(), causing an + IndexError when the data are not time-ordered. + """ + import numpy as np + idx = np.argsort(transient_object.x) + for attr in ('x', 'y', 'y_err', 'x_err', 'bands', 'frequencies'): + arr = getattr(transient_object, attr, None) + if arr is not None and hasattr(arr, '__len__') and len(arr) == len(idx): + setattr(transient_object, attr, arr[idx]) + + +def fit_redback_model( + transient_object, + object_name, + model='tde_analytical', + sampler='dynesty', + nlive=500, + npool=4, + resume=True, + prior_overrides=None, + multiband_filters=None, + random_models=100, + outdir_root='redback_inference', +): + """ + Run redback inference on a transient object and plot results. + + Parameters + ---------- + transient_object : redback transient + Output of create_redback_transient. + object_name : str + Transient name — used for outdir and result label. + model : str + Redback model name. Default: 'tde_analytical'. + sampler : str + Bilby sampler. Default: 'dynesty'. + nlive : int + Number of live points. Default: 500. + npool : int + Number of parallel processes. Default: 4. + resume : bool + Resume from existing checkpoint if available. Default: True. + prior_overrides : dict or None + Parameter overrides applied on top of the default priors. + Values can be floats (fixed) or bilby Prior objects. + Example:: + + import bilby + prior_overrides = { + 'redshift': 0.403, + 'temperature_0': bilby.core.prior.Uniform(10000, 60000, + name='temperature_0', + latex_label='$T_0$ (K)'), + } + + multiband_filters : list of str or None + Filters for plot_multiband_lightcurve (e.g. ['ztfg', 'ztfr']). + If None, the multiband plot is skipped. + random_models : int + Number of posterior draws to overplot. Default: 100. + outdir_root : str + Root directory for inference outputs. Default: 'redback_inference'. + + Returns + ------- + result : redback result + Bilby/redback result object. + """ + # Workaround for redback bug in Diffusion.convert_input_luminosity: + # dense_times uses time[-1] as the max, which is only correct if times are + # sorted. Unsorted data causes searchsorted to return an out-of-bounds index. + _sort_transient_by_time(transient_object) + + outdir = f'{outdir_root}/{object_name}/{model}' + os.makedirs(outdir, exist_ok=True) + + priors = redback.priors.get_priors(model=model) + if prior_overrides: + for key, val in prior_overrides.items(): + priors[key] = val + + model_kwargs = dict( + frequency=transient_object.filtered_frequencies, + output_format='flux_density', + ) + + result = redback.fit_model( + transient=transient_object, + model=model, + sampler=sampler, + model_kwargs=model_kwargs, + prior=priors, + sample='rslice', + nlive=nlive, + npool=npool, + resume=resume, + outdir=outdir, + label=f'{object_name}_nlive{nlive}', + plot=False, + ) + + # result.transient reconstructs from metadata with wrong path; + # use transient_object directly for plots + model_func = redback.model_library.all_models_dict[model] + + result.plot_corner(save=False, show=True) + plt.show() + plt.close('all') + + transient_object.plot_lightcurve( + model=model_func, + posterior=result.posterior, + model_kwargs=result.model_kwargs, + random_models=random_models, + save=False, + show=True, + ) + + if multiband_filters is not None: + transient_object.plot_multiband_lightcurve( + model=model_func, + posterior=result.posterior, + model_kwargs=result.model_kwargs, + random_models=random_models, + filters=multiband_filters, + save=False, + show=True, + ) + + return result