Add Result Unraveling Function (#1047)

* initial

* black formatting

* address comments on pr #1047

Author: maxwest-uw

src/kbmod/util_functions.py

@@ -6,6 +6,7 @@
 from astropy.io import fits
 from astropy.time import Time
 from itertools import product
+import pandas as pd
 
 from kbmod.core.image_stack_py import LayeredImagePy
 
@@ -114,3 +115,121 @@ def load_deccam_layered_image(filename, psf):
         )
 
     return img
+
+
+def get_unique_obstimes(all_obstimes):
+    """Get the unique observation times and their indices.
+    Used to group observations for mosaicking.
+
+    Parameters
+    ----------
+    all_obstimes : `np.ndarray`
+        The array of observation times.
+
+    Returns
+    -------
+    unique_obstimes : `np.ndarray`
+        The unique observation times.
+    unique_indices : `list`
+        A list of lists, where each sublist contains the indices of the grouping.
+    """
+    unique_obstimes = np.unique(all_obstimes)
+    unique_indices = [list(np.where(all_obstimes == time)[0]) for time in unique_obstimes]
+    return unique_obstimes, unique_indices
+
+
+def get_magnitude(flux, zero_point):
+    """Convert a flux value to a magnitude using the zero point.
+
+    Parameters
+    ----------
+    flux : `float`
+        The flux value to convert.
+    zero_point : `float`
+        The zero point of the observations.
+
+    Returns
+    -------
+    mag : `float`
+        The calculated magnitude.
+    """
+    mag = -2.5 * np.log10(flux) + zero_point
+    return mag
+
+
+def unravel_results(results, image_collection, obscode="X05", batch_id=None):
+    """Take a results file and transform it into a table of individual observations.
+
+    Parameters
+    ----------
+    results : `kbmod.results.Results`
+        The results.
+    image_collection : `kbmod.image_collection.ImageCollection`
+        The image collection containing the images used in the results.
+    obscode : `str`, optional
+        The observatory code to use for the observations.
+        Default: "X05" (LSST).
+    batch_id : `str`, optional
+        The batch ID to use for this result set.
+        individual observation ids will be in the format of
+        "{batch_id}-{result #}-{observation #}".
+
+    Returns
+    -------
+    final_df : `pandas.DataFrame`
+        A DataFrame containing the individual observations with columns:
+        - id: The unique identifier for the observation.
+        - ra: The right ascension of the observation in degrees.
+        - dec: The declination of the observation in degrees.
+        - magnitude: The magnitude of the observation.
+        - mjd: The modified Julian date of the observation.
+        - band: The band of the observation.
+        - obscode: The observatory code for the observation.
+    """
+    zp = np.mean(image_collection["zeroPoint"])
+
+    ids = []
+    ras = []
+    decs = []
+    mags = []
+    mjds = []
+    bands = []
+    obscodes = []
+
+    all_times = results.table.meta["mjd_mid"]
+    all_bands = image_collection["band"]
+
+    _, unique_indices = get_unique_obstimes(image_collection["mjd_mid"])
+    first_of_each_frame = np.array([i[0] for i in unique_indices])
+
+    for i, row in enumerate(results):
+        if "obs_valid" in results.table.colnames:
+            valid_obs = row["obs_valid"]
+        else:
+            valid_obs = np.full(row["obs_count"], True)
+        num_valid = row["obs_count"]
+
+        # need to figure out a better way to do this
+        if batch_id is not None:
+            ids.append([f"{batch_id}-{i}-{j}" for j in range(num_valid)])
+        else:
+            ids.append([f"{i}-{j}" for j in range(num_valid)])
+
+        ras.append(row["img_ra"][valid_obs])
+        decs.append(row["img_dec"][valid_obs])
+
+        mags.append([get_magnitude(row["flux"], zp)] * num_valid)
+        mjds.append(all_times[valid_obs])
+        bands.append(all_bands[first_of_each_frame][valid_obs])
+        obscodes.append([obscode] * num_valid)
+
+    final_df = pd.DataFrame()
+    final_df["id"] = np.concatenate(ids)
+    final_df["ra"] = np.concatenate(ras)
+    final_df["dec"] = np.concatenate(decs)
+    final_df["magnitude"] = np.concatenate(mags)
+    final_df["mjd"] = np.concatenate(mjds)
+    final_df["band"] = np.concatenate(bands)
+    final_df["obscode"] = np.concatenate(obscodes)
+
+    return final_df

tests/test_util_functions.py

@@ -1,14 +1,18 @@
+from astropy.table import Table
 import numpy as np
 import unittest
 
 from pathlib import Path
 
 from kbmod.core.image_stack_py import LayeredImagePy
 from kbmod.core.psf import PSF
+from kbmod.results import Results
+from kbmod.search import Trajectory
 from kbmod.util_functions import (
     get_matched_obstimes,
     load_deccam_layered_image,
     mjd_to_day,
+    unravel_results,
 )
 
 
@@ -34,6 +38,48 @@ def test_load_deccam_layered_image(self):
         self.assertGreater(img.width, 0)
         self.assertGreater(img.height, 0)
 
+    def test_unravel_results(self):
+        trj_list = []
+        num_images = 10
+        num_trjs = 11
+
+        for i in range(num_trjs):
+            trj = Trajectory(
+                x=i,
+                y=i + 0,
+                vx=i - 2.0,
+                vy=i + 5.0,
+                flux=5.0 * i,
+                lh=100.0 + i,
+                obs_count=num_images,
+            )
+            trj_list.append(trj)
+
+        res = Results.from_trajectories(trj_list)
+
+        # create an "image collection" with all required fields
+        ic = Table()
+        ic["mjd_mid"] = [60000 + i for i in range(num_images)]
+        ic["band"] = ["g" for _ in range(num_images)]
+        ic["zeroPoint"] = [31.4 for _ in range(num_images)]
+
+        res.table.meta["mjd_mid"] = ic["mjd_mid"]
+        obs_count = [num_images for _ in range(num_trjs)]
+        res.table["obs_count"] = obs_count
+        res.table["img_ra"] = [np.array([j + (i * 0.1) for j in range(num_images)]) for i in range(num_trjs)]
+        res.table["img_dec"] = [np.array([j + (i * 0.1) for j in range(num_images)]) for i in range(num_trjs)]
+
+        df = unravel_results(res, ic)
+        self.assertEqual(len(df), (num_images * num_trjs))
+
+        obs_count[int(num_images / 2)] = num_images - 1
+        valids = [[True] * num_images for _ in range(num_trjs)]
+        valids[int(num_images / 2)][-1] = False  # make one observation invalid
+        res.table["obs_valid"] = valids
+        res.table["obs_count"] = obs_count
+        df2 = unravel_results(res, ic)
+        self.assertEqual(len(df2), (num_images * num_trjs) - 1)
+
 
 if __name__ == "__main__":
     unittest.main()