From 817335ea12be04a3094a1989448fd7717e45578c Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 13 Aug 2026 15:30:20 +1000
Subject: [PATCH 01/11] Split output of parallel_groupby into coords (float
 vars) and lists of strings

---
 vast_pipeline/pipeline/utils.py | 63 ++++++++++++++++++++-------------
 1 file changed, 39 insertions(+), 24 deletions(-)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index b24ef05ef..13696f744 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -760,9 +760,19 @@ def calc_ave_coord(grp: pd.DataFrame) -> pd.Series:
     return pd.Series(d)
 
 
-def parallel_groupby_coord(df: dd.DataFrame,) -> pd.DataFrame:
+def parallel_groupby_coord(df: dd.DataFrame,) -> Tuple[pd.DataFrame, pd.DataFrame]:
     """Calculate the weighted average RA and Dec of the sources.
 
+    Produces two separate per-source DataFrames in a single Dask compute pass:
+
+    * **coords_df** — lightweight numeric frame (one float per column) used
+      for the AstroPy sky crossmatch: ``wavg_ra``, ``wavg_dec``,
+      ``flux_peak``.
+    * **lists_df** — heavyweight frame holding Python list columns
+      ``img_list`` and ``epoch_list``, only needed for the "missing image"
+      computation in ``get_src_skyregion_merged_df``.  Keeping it separate
+      means the large list objects are not in memory during the crossmatch.
+
     NOTE: Sergio had the idea to persist the dataframe result and keep it in the
     cluster. However since then the ideal image method uses the astropy match sky
     method which relies on being able to iloc the dataframe. This would be really
@@ -772,31 +782,36 @@ def parallel_groupby_coord(df: dd.DataFrame,) -> pd.DataFrame:
         df: The sources dataframe.
 
     Returns:
-        The resulting average coordinate values and unique image and epoch
-            lists for each unique source (group).
+        Tuple of (coords_df, lists_df) — both indexed by source id.
     """
-    cols = [
-        'source', 'image', 'epoch', 'interim_ew', 'weight_ew', 'interim_ns', 'weight_ns'
+
+    coord_cols = [
+        'source', 'interim_ew', 'weight_ew', 'interim_ns', 'weight_ns',
     ]
-    cols_to_sum = ['interim_ew', 'weight_ew', 'interim_ns', 'weight_ns']
-    aggregations = {'interim_ew': 'sum',
-                    'weight_ew': 'sum',
-                    'interim_ns': 'sum',
-                    'weight_ns': 'sum',
-                    'image': list,
-                    'epoch': list}
-
-    groups = df[cols].groupby('source')
-    out = groups.agg(aggregations)
-    out['wavg_ra'] = out['interim_ew'] / out['weight_ew']
-    out['wavg_dec'] = out['interim_ns'] / out['weight_ns']
-    out = out.drop(cols_to_sum, axis=1).rename(columns={'image': 'img_list', 'epoch': 'epoch_list'})
-
-    # Do the aggregations now.
-    out = out.compute()
-
-    del groups
-    return out
+    coord_agg = {
+        'interim_ew': 'sum',
+        'weight_ew': 'sum',
+        'interim_ns': 'sum',
+        'weight_ns': 'sum',
+    }
+    coord_groups = df[coord_cols].groupby('source').agg(coord_agg)
+
+    list_cols = ['source', 'image', 'epoch']
+    list_agg = {'image': list, 'epoch': list}
+    list_groups = df[list_cols].groupby('source').agg(list_agg)
+
+    coords_raw, lists_raw = dask.compute(coord_groups, list_groups)
+
+    coords_df = coords_raw
+    coords_df['wavg_ra'] = coords_df['interim_ew'] / coords_df['weight_ew']
+    coords_df['wavg_dec'] = coords_df['interim_ns'] / coords_df['weight_ns']
+    coords_df = coords_df.drop(
+        columns=['interim_ew', 'weight_ew', 'interim_ns', 'weight_ns']
+    )
+
+    lists_df = lists_raw.rename(columns={'image': 'img_list', 'epoch': 'epoch_list'})
+
+    return coords_df, lists_df
 
 
 def get_rms_noise_image_values(rms_path: str) -> Tuple[float, float, float]:

From 78ed4ac7051648c23b19c71b3a5e3deab2b92af0 Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 13 Aug 2026 15:45:29 +1000
Subject: [PATCH 02/11] Add _build_compact_indices function to convert strings
 to int32 lookups and conserve memory

---
 vast_pipeline/pipeline/utils.py | 57 +++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index 13696f744..e1b1bb397 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -848,6 +848,63 @@ def get_rms_noise_image_values(rms_path: str) -> Tuple[float, float, float]:
     return med_val, min_val, max_val
 
 
+def _build_compact_indices(
+    sources_df: dd.DataFrame, images_df: pd.DataFrame,
+) -> Tuple[dd.DataFrame, np.ndarray, int, pd.DataFrame, int]:
+    """
+    Replaces image name strings with compact int32 codes for the rest of
+    `get_src_skyregion_merged_df`, and builds the per-image ideal-coverage
+    frame used by the later crossmatch.
+
+    Args:
+        sources_df: The association step output, must have an 'image'
+            column holding image name strings.
+        images_df: All image objects for the run, with 'name', 'skyreg_id',
+            'epoch' and 'datetime' columns.
+
+    Returns:
+        sources_df: With the 'image' column replaced by int32 codes.
+        image_names: Array mapping int32 image code -> original image name.
+        img_mult: Multiplier for the combined (source, image) key used by
+            the vectorized "missing image" membership test.
+        skyreg_img_df: Per-image ideal-coverage frame indexed by
+            'skyreg_id', columns 'skyreg_img_list' (int32 image code),
+            'skyreg_epoch' (int32) and 'skyreg_datetime' (int64, a sort key
+            only, not a real datetime).
+        epoch_mult: Multiplier for the combined (source, epoch) key.
+    """
+    # Compact int32 image code replaces the name string from here on;
+    # image_names converts back to names in get_src_skyregion_merged_df.
+    image_names = images_df["name"].to_numpy()
+    name_to_idx = pd.Series(np.arange(len(image_names), dtype=np.int32), index=image_names)
+    images_df = images_df.assign(name=np.arange(len(image_names), dtype=np.int32))
+    sources_df["image"] = sources_df["image"].map(name_to_idx, meta=("image", "int32"))
+    # Multiplier for the combined (source, image) key used by
+    # _compute_missing_images' vectorized membership test.
+    img_mult = len(image_names)
+
+    skyreg_img_df = images_df[["skyreg_id", "name", "epoch", "datetime"]].rename(
+        columns={
+            "name": "skyreg_img_list",
+            "epoch": "skyreg_epoch",
+            "datetime": "skyreg_datetime",
+        }
+    )
+    # int32/int64 downcasts halve these columns' cost across the large
+    # crossmatch expansion in _crossmatch_sources_to_skyregions.
+    skyreg_img_df["skyreg_epoch"] = skyreg_img_df["skyreg_epoch"].astype(np.int32)
+    # Multiplier for the combined (source, epoch) key used by
+    # _compute_missing_images; captured now while skyreg_img_df still holds
+    # the full epoch universe.
+    epoch_mult = int(skyreg_img_df["skyreg_epoch"].max()) + 1
+    # Avoids boxing tz-aware Timestamps in the merges/sorts that follow;
+    # only used to establish sort order, then dropped.
+    skyreg_img_df["skyreg_datetime"] = skyreg_img_df["skyreg_datetime"].astype(np.int64)
+    skyreg_img_df = skyreg_img_df.set_index("skyreg_id")
+
+    return sources_df, image_names, img_mult, skyreg_img_df, epoch_mult
+
+
 def get_image_list_diff(row: pd.Series) -> Union[List[str], int]:
     """
     Calculate the difference between the ideal coverage image list of a source

From 8520c4537dc22ad46eb7ce7d852643564b4460dd Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 13 Aug 2026 16:00:45 +1000
Subject: [PATCH 03/11] Move crossmatching to separate function.

---
 vast_pipeline/pipeline/utils.py | 60 +++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index e1b1bb397..9e9bbbac1 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -6,6 +6,7 @@
 import os
 import logging
 import glob
+import gc
 import shutil
 import numpy as np
 import pandas as pd
@@ -905,6 +906,65 @@ def _build_compact_indices(
     return sources_df, image_names, img_mult, skyreg_img_df, epoch_mult
 
 
+def _crossmatch_sources_to_skyregions(
+    coords_df: pd.DataFrame, skyreg_df: pd.DataFrame, skyreg_img_df: pd.DataFrame,
+) -> pd.DataFrame:
+    """
+    Crossmatches each source with every sky region within that region's
+    extraction radius, then expands each match with every image belonging
+    to the sky region -- i.e. the source's "ideal coverage" images/epochs.
+
+    Args:
+        coords_df: Per-source average coordinates, indexed by int32 source
+            code. Only 'wavg_ra'/'wavg_dec' are used.
+        skyreg_df: Sky regions of the run, with 'id', 'centre_ra',
+            'centre_dec' and 'xtr_radius' columns.
+        skyreg_img_df: Per-image ideal-coverage frame indexed by
+            'skyreg_id', as returned by `_build_compact_indices`.
+
+    Returns:
+        Dataframe with one row per (source, sky region, ideal image) match
+        and columns 'source', 'sep', 'skyreg_img_list', 'skyreg_epoch' and
+        'skyreg_datetime'.
+    """
+    skyreg_df = skyreg_df[["id", "centre_ra", "centre_dec", "xtr_radius"]]
+
+    # crossmatch sources with sky regions up to the max sky region radius
+    skyreg_coords = SkyCoord(
+        ra=skyreg_df.centre_ra.values, dec=skyreg_df.centre_dec.values, unit="deg"
+    )
+    srcs_coords = SkyCoord(
+        ra=coords_df["wavg_ra"],
+        dec=coords_df["wavg_dec"],
+        unit="deg")
+    skyreg_idx, srcs_idx, sep, _ = srcs_coords.search_around_sky(
+        skyreg_coords, skyreg_df.xtr_radius.values * u.deg
+    )
+    skyreg_df = skyreg_df.drop(
+        columns=[
+            "centre_ra",
+            "centre_dec",
+            "xtr_radius"]).set_index("id")
+
+    # Build the per-source ideal-images frame
+    src_skyrg_df = pd.DataFrame(
+        {
+            "source": coords_df.iloc[srcs_idx].index,
+            "sep": sep.to("deg").value.astype(np.float32),
+        },
+        index=skyreg_df.iloc[skyreg_idx].index,
+    )
+
+    src_skyrg_df = src_skyrg_df.join(skyreg_df, how="inner")
+    src_skyrg_df = src_skyrg_df.join(skyreg_img_df, how="inner")
+
+    src_skyrg_df = src_skyrg_df.reset_index(drop=True)
+
+    del skyreg_df, skyreg_img_df
+    gc.collect()
+
+    return src_skyrg_df
+
 def get_image_list_diff(row: pd.Series) -> Union[List[str], int]:
     """
     Calculate the difference between the ideal coverage image list of a source

From b624c9682dbd5f15febcc4b5b2d87c5f5dbd47f0 Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 13 Aug 2026 16:50:44 +1000
Subject: [PATCH 04/11] Add vectorised and numpy based operations for getting
 missing images.

---
 vast_pipeline/pipeline/main.py  |  16 +-
 vast_pipeline/pipeline/utils.py | 362 ++++++++++++++++++++------------
 2 files changed, 231 insertions(+), 147 deletions(-)

diff --git a/vast_pipeline/pipeline/main.py b/vast_pipeline/pipeline/main.py
index 050e4f3f4..8014515ff 100644
--- a/vast_pipeline/pipeline/main.py
+++ b/vast_pipeline/pipeline/main.py
@@ -285,14 +285,14 @@ def process_pipeline(self, p_run: Run) -> None:
         # Make missing sources into Dask dataframe
         # NOTE: This would not be necessary if the get_src_skyregion_merged_df
         # function was improved to use Dask. (See NOTE in parallel_groupby function.)
-        npartitions = calculate_n_partitions(missing_sources_df,
-                                             n_cpu=self.dm.num_workers,
-                                             partition_size_mb=self.config['processing']['max_partition_mb'])
-        missing_sources_df = dd.from_pandas(
-            missing_sources_df,
-            npartitions=npartitions
-        )
-        wait(missing_sources_df)
+        #npartitions = calculate_n_partitions(missing_sources_df,
+        #                                     n_cpu=self.dm.num_workers,
+        #                                     partition_size_mb=self.config['processing']['max_partition_mb'])
+        #missing_sources_df = dd.from_pandas(
+        #    missing_sources_df,
+        #    npartitions=npartitions
+        #)
+        #wait(missing_sources_df)
 
         # STEP #4 New source analysis
         logger.info("Running step #4: new source analysis...")
diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index 9e9bbbac1..ab8477851 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -965,84 +965,215 @@ def _crossmatch_sources_to_skyregions(
 
     return src_skyrg_df
 
-def get_image_list_diff(row: pd.Series) -> Union[List[str], int]:
+
+def _dedupe_closest_skyregion_per_epoch(src_skyrg_df: pd.DataFrame) -> pd.DataFrame:
     """
-    Calculate the difference between the ideal coverage image list of a source
-    and the actual observed image list. Also checks whether an epoch does in
-    fact contain a detection but is not in the expected 'ideal' image for that
-    epoch.
+    For each (source, ideal epoch) pair, keeps only the closest-matching
+    sky region (smallest separation), then sorts the result ready for the
+    vectorized "missing image" computation that follows.
 
     Args:
-        row: The row from the sources dataframe that is being iterated over.
+        src_skyrg_df: Output of `_crossmatch_sources_to_skyregions`, one
+            row per (source, sky region, ideal image) match, still with
+            'sep' and 'skyreg_datetime' columns.
 
     Returns:
-        A list of the images missing from the observed image list.
-    """
-    out = list(filter(lambda arg: arg not in row["img_list"], row["skyreg_img_list"]))
+        One row per (source, ideal image), sorted by 'source' (stable
+        sort, so each source's rows stay in the chronological order
+        established via 'skyreg_datetime'), with 'sep' and
+        'skyreg_datetime' dropped.
+    """
+    # Sort by (source, skyreg_epoch, sep) so every row sharing a
+    # (source, skyreg_epoch) pair is contiguous, smallest sep first.
+    src_skyrg_df.sort_values(
+        ['source', 'skyreg_epoch', 'sep'], inplace=True
+    )
 
-    # Check that an epoch has not already been seen (just not in the 'ideal'
-    # image)
-    out_epochs = [
-        row["skyreg_epoch"][pair[0]]
-        for pair in enumerate(row["skyreg_img_list"])
-        if pair[1] in out
-    ]
+    # Use numpy arrays to find the first row of each (source, skyreg_epoch) group
+    # and drop the rest (keeping only the closest sky region per ideal epoch).
+    source_arr = src_skyrg_df["source"].to_numpy()
+    epoch_arr = src_skyrg_df["skyreg_epoch"].to_numpy()
+    is_first = np.empty(len(source_arr), dtype=bool)
+    is_first[0] = True
+    is_first[1:] = (source_arr[1:] != source_arr[:-1]) | (epoch_arr[1:] != epoch_arr[:-1])
 
-    out = [
-        out[pair[0]]
-        for pair in enumerate(out_epochs)
-        if pair[1] not in row["epoch_list"]
-    ]
+    del source_arr, epoch_arr
 
-    return out
+    src_skyrg_df = src_skyrg_df[is_first].drop(columns=["sep"])
+
+    # Now sort by datetime int64 value
+    src_skyrg_df.sort_values(by="skyreg_datetime", inplace=True)
+    src_skyrg_df.drop(columns=["skyreg_datetime"], inplace=True)
+
+    # Stable sort to preserve the chronological order of each source's rows
+    src_skyrg_df.sort_values(by="source", kind="stable", inplace=True)
+
+    return src_skyrg_df
 
 
-def get_names_and_epochs(grp: pd.DataFrame) -> pd.Series:
+def _isin_sorted(keys: np.ndarray, sorted_unique_ref: np.ndarray) -> np.ndarray:
     """
-    Convenience function to group together the image names, epochs and
-    datetimes into one list object which is then returned as a pandas series.
-    This is necessary for easier processing in the ideal coverage analysis.
+    Vectorized membership test of `keys` against a sorted, unique reference
+    array, via binary search.
 
     Args:
-        grp: A group from the grouped by sources DataFrame.
+        keys: Array of int64 keys to test.
+        sorted_unique_ref: Sorted, unique int64 reference array.
 
     Returns:
-        Pandas series containing the list object that contains the lists of the
-            image names, epochs and datetimes.
+        Boolean array, True where the corresponding key is present in
+        `sorted_unique_ref`.
     """
-    d = {}
-    d["skyreg_img_epoch_list"] = [
-        [
-            [
-                x,
-            ],
-            y,
-            z,
-        ]
-        for x, y, z in zip(
-            grp["name"].values.tolist(),
-            grp["epoch"].values.tolist(),
-            grp["datetime"].values.tolist(),
-        )
-    ]
+    idx = np.searchsorted(sorted_unique_ref, keys)
+    idx = np.clip(idx, 0, len(sorted_unique_ref) - 1)
+    return (idx < len(sorted_unique_ref)) & (sorted_unique_ref[idx] == keys)
 
-    return pd.Series(d)
 
+def _compute_missing_images(
+    src_skyrg_df: pd.DataFrame, lists_df: pd.DataFrame, img_mult: int, epoch_mult: int,
+) -> Tuple[pd.Series, pd.DataFrame]:
+    """
+    For each source, determines which ideal images/epochs were never
+    actually observed, plus the source's first ideal ("primary") and first
+    observed ("detection") image. Uses vectorized lookups.
 
-def check_primary_image(row: pd.Series) -> bool:
+    Args:
+        src_skyrg_df: Output of `_dedupe_closest_skyregion_per_epoch`, one
+            row per (source, ideal image), sorted by 'source'.
+        lists_df: Per-source 'img_list'/'epoch_list' columns (the actually
+            observed images/epochs), indexed by int32 source code.
+        img_mult: Multiplier for the combined (source, image) key.
+        epoch_mult: Multiplier for the combined (source, epoch) key.
+
+    Returns:
+        img_diff_series: Per-source list of missing ideal images, indexed
+            by source, name 'img_diff'. Only sources with >=1 missing
+            image are present.
+        per_source_df: Per-source 'detection' (first observed image) and
+            'in_primary' (whether the source was detected in its first
+            ideal image) columns, indexed by source.
+    """
+    # For each (source, ideal image)/(source, ideal epoch) pair, test
+    # whether it was ever actually observed by encoding the pair as one
+    # combined int64 key (source * multiplier + value) and binary-searching
+    # it (np.searchsorted) against a sorted array of observed keys.
+    source_arr = src_skyrg_df["source"].to_numpy()
+    skyreg_img_arr = src_skyrg_df["skyreg_img_list"].to_numpy()
+    skyreg_epoch_arr = src_skyrg_df["skyreg_epoch"].to_numpy()
+    del src_skyrg_df
+
+    split_points = np.flatnonzero(np.diff(source_arr)) + 1
+    group_start_idx = np.concatenate(([0], split_points))
+    group_source = source_arr[group_start_idx]
+    # "primary" = first (chronologically earliest) ideal image per source.
+    primary_arr = skyreg_img_arr[group_start_idx]
+    del group_start_idx
+
+    # Long-format "observed" (source, image)/(source, epoch) pairs, built by
+    # exploding the small (one row per source) img_list/epoch_list columns.
+    obs_img_long = lists_df["img_list"].explode()
+    obs_epoch_long = lists_df["epoch_list"].explode()
+    obs_img_keys = np.unique(
+        obs_img_long.index.to_numpy().astype(np.int64) * img_mult
+        + obs_img_long.to_numpy().astype(np.int64)
+    )
+    obs_epoch_keys = np.unique(
+        obs_epoch_long.index.to_numpy().astype(np.int64) * epoch_mult
+        + obs_epoch_long.to_numpy().astype(np.int64)
+    )
+    del obs_img_long, obs_epoch_long
+
+    ideal_img_key = source_arr.astype(np.int64) * img_mult + skyreg_img_arr.astype(np.int64)
+    in_img_list = _isin_sorted(ideal_img_key, obs_img_keys)
+    del ideal_img_key
+
+    ideal_epoch_key = source_arr.astype(np.int64) * epoch_mult + skyreg_epoch_arr.astype(np.int64)
+    in_epoch_list = _isin_sorted(ideal_epoch_key, obs_epoch_keys)
+    del ideal_epoch_key, skyreg_epoch_arr
+
+    missing_mask = ~in_img_list & ~in_epoch_list
+    del in_img_list, in_epoch_list
+
+    missing_source = source_arr[missing_mask]
+    missing_img = skyreg_img_arr[missing_mask]
+    del skyreg_img_arr, missing_mask
+
+    # missing_source is a subset of the already (stable-)sorted source_arr,
+    # so it's still sorted — group it straight back into per-source lists
+    # without needing to re-sort.
+    if len(missing_source) > 0:
+        m_split_points = np.flatnonzero(np.diff(missing_source)) + 1
+        m_group_source = missing_source[np.concatenate(([0], m_split_points))]
+        img_diff_groups = np.split(missing_img, m_split_points)
+    else:
+        m_group_source = np.array([], dtype=source_arr.dtype)
+        img_diff_groups = []
+    del missing_source, missing_img, source_arr
+
+    img_diff_series = pd.Series(
+        img_diff_groups, index=pd.Index(m_group_source, name="source"), name="img_diff",
+    )
+    del m_group_source, img_diff_groups
+
+    # "detection" = first (chronologically earliest) *observed* image per source.
+    detection_series = lists_df["img_list"].str[0]
+    primary_key = group_source.astype(np.int64) * img_mult + primary_arr.astype(np.int64)
+    in_primary_arr = _isin_sorted(primary_key, obs_img_keys)
+    del primary_key, obs_img_keys, obs_epoch_keys, primary_arr
+
+    per_source_df = pd.DataFrame(
+        {
+            "detection": detection_series,
+            "in_primary": pd.Series(in_primary_arr, index=pd.Index(group_source, name="source")),
+        }
+    )
+    del detection_series, in_primary_arr, group_source
+
+    return img_diff_series, per_source_df
+
+
+def _explode_missing_images_to_dask(
+    srcs_df: pd.DataFrame, image_names: np.ndarray, n_cpu: Optional[int],
+) -> dd.DataFrame:
     """
-    Checks whether the primary image of the ideal source
-    dataframe is in the image list for the source.
+    Converts the assembled per-source result to a Dask DataFrame and
+    explodes 'img_diff' to one row per missing image, pre-computing the
+    exploded form shared by the pipeline's steps #4 and #5.
 
     Args:
-        row:
-            Input dataframe row, with columns ['primary'] and ['img_list'].
+        srcs_df: One row per source, with 'wavg_ra', 'wavg_dec',
+            'flux_peak', 'detection' (image name), 'in_primary' and
+            'img_diff' (a list of int32 image codes), indexed by source
+            ShortUUID.
+        image_names: Array mapping int32 image code -> original image name.
+        n_cpu: Number of available CPUs/workers, used to size the returned
+            Dask DataFrame's partitions (via `calculate_n_partitions`). If
+            None, partitions are sized on memory alone.
 
     Returns:
-        True if primary in image list else False.
+        Dask DataFrame with one row per (source, missing image) pair, as
+        described in `get_src_skyregion_merged_df`.
     """
-    return row["primary"] in row["img_list"]
+    base_npartitions = calculate_n_partitions(srcs_df, n_cpu=n_cpu, partition_size_mb=10)
+    exploded_npartitions = max(n_cpu, base_npartitions) if n_cpu else base_npartitions
+
+    srcs_df = dd.from_pandas(srcs_df, npartitions=exploded_npartitions)
+
+    srcs_df = srcs_df.reset_index()[
+        ["source", "wavg_ra", "wavg_dec", "img_diff", "detection", "in_primary", "flux_peak"]
+    ].explode("img_diff")
+
+    def _convert_img_diff_names(partition: pd.DataFrame) -> pd.DataFrame:
+        """Vectorized int32-code -> image-name lookup for one partition"""
+        partition = partition.copy()
+        partition["img_diff"] = image_names[partition["img_diff"].to_numpy().astype(np.int32)]
+        return partition
+
+    srcs_df = srcs_df.map_partitions(
+        _convert_img_diff_names, meta=srcs_df._meta.assign(img_diff=pd.Series(dtype=object)),
+    )
 
+    return srcs_df
 
 def get_src_skyregion_merged_df(
     sources_df: dd.DataFrame, images_df: pd.DataFrame, skyreg_df: pd.DataFrame
@@ -1108,103 +1239,56 @@ def get_src_skyregion_merged_df(
 
     merged_timer = StopWatch()
 
-    skyreg_df = skyreg_df.drop(["x", "y", "z", "width_ra", "width_dec"], axis=1)
-
-    skyreg_df = skyreg_df.join(
-        pd.DataFrame(images_df.groupby("skyreg_id")[["skyreg_id", "name", "epoch", "datetime"]]
-                              .apply(get_names_and_epochs)),
-        on="id",
+    sources_df, image_names, img_mult, skyreg_img_df, epoch_mult = _build_compact_indices(
+        sources_df, images_df
     )
 
-    # calculate some metrics on sources
-    # compute only some necessary metrics in the groupby
-    timer = StopWatch()
-    srcs_df = parallel_groupby_coord(sources_df)
-    logger.debug('Groupby-apply time: %.2f seconds', timer.reset())
 
+    coords_df, lists_df = parallel_groupby_coord(sources_df)
+    # coords_df: wavg_ra/wavg_dec/flux_peak, used by _crossmatch_sources_to_skyregions.
+    # lists_df: img_list/epoch_list, used by _compute_missing_images.
     del sources_df
 
-    # crossmatch sources with sky regions up to the max sky region radius
-    skyreg_coords = SkyCoord(
-        ra=skyreg_df.centre_ra, dec=skyreg_df.centre_dec, unit="deg"
-    )
-    srcs_coords = SkyCoord(
-        ra=srcs_df["wavg_ra"],
-        dec=srcs_df["wavg_dec"],
-        unit="deg")
-    skyreg_idx, srcs_idx, sep, _ = srcs_coords.search_around_sky(
-        skyreg_coords, skyreg_df.xtr_radius.max() * u.deg
-    )
-    skyreg_df = skyreg_df.drop(
-        columns=[
-            "centre_ra",
-            "centre_dec"]).set_index("id")
-
-    # select rows where separation is less than sky region radius
-    # drop not more useful columns and groupby source id
-    # compute list of images
-    src_skyrg_df = (
-        pd.DataFrame(
-            {
-                "source": srcs_df.iloc[srcs_idx].index,
-                "id": skyreg_df.iloc[skyreg_idx].index,
-                "sep": sep.to("deg").value,
-            }
-        )
-        .merge(skyreg_df, left_on="id", right_index=True)
-        .query("sep < xtr_radius")
-        .drop(columns=["id", "xtr_radius"])
-        .explode("skyreg_img_epoch_list")
-    )
+    # Use compact int32 source code instead of the ShortUUID string for the
+    # rest of this function; converted back on the final index only.
+    source_ids = coords_df.index.to_numpy()
+    source_to_idx = pd.Series(np.arange(len(source_ids), dtype=np.int32), index=source_ids)
+    coords_df.index = np.arange(len(source_ids), dtype=np.int32)
+    coords_df.index.name = "source"
+    lists_df.index = lists_df.index.map(source_to_idx)
+    lists_df.index.name = "source"
 
+    # crossmatch sources with sky regions up to the max sky region radius
+    src_skyrg_df = _crossmatch_sources_to_skyregions(coords_df, skyreg_df, skyreg_img_df)
     del skyreg_df
 
-    src_skyrg_df[["skyreg_img_list", "skyreg_epoch", "skyreg_datetime"]] = pd.DataFrame(
-        src_skyrg_df["skyreg_img_epoch_list"].tolist(), index=src_skyrg_df.index
-    )
-
-    src_skyrg_df = src_skyrg_df.drop("skyreg_img_epoch_list", axis=1)
-
-    src_skyrg_df = (
-        src_skyrg_df.sort_values(["source", "sep"])
-        .drop_duplicates(["source", "skyreg_epoch"])
-        .sort_values(by="skyreg_datetime")
-        .drop(["sep", "skyreg_datetime"], axis=1)
-    )
-    # annoyingly epoch needs to be not a list to drop duplicates
-    # but then we need to sum the epochs into a list
-    src_skyrg_df["skyreg_epoch"] = src_skyrg_df["skyreg_epoch"].apply(
-        lambda x: [
-            x,
-        ]
-    )
-
-    src_skyrg_df = src_skyrg_df.groupby("source").sum(
-        numeric_only=False
-    )  # sum because we need to preserve order
+    # drop duplicates of the same source and sky region for the same epoch, keeping only the closest match
+    src_skyrg_df = _dedupe_closest_skyregion_per_epoch(src_skyrg_df)
 
-    # merge into main df and compare the images
-    srcs_df = srcs_df.merge(src_skyrg_df, left_index=True, right_index=True)
-
-    del src_skyrg_df
-
-    srcs_df["img_diff"] = srcs_df[
-        ["img_list", "skyreg_img_list", "epoch_list", "skyreg_epoch"]
-    ].apply(get_image_list_diff, axis=1)
-
-    srcs_df = srcs_df.loc[srcs_df["img_diff"].apply(len) > 0]
-
-    srcs_df = srcs_df.drop(["epoch_list", "skyreg_epoch"], axis=1)
-
-    srcs_df["primary"] = srcs_df["skyreg_img_list"].apply(lambda x: x[0])
-
-    srcs_df["detection"] = srcs_df["img_list"].apply(lambda x: x[0])
-
-    srcs_df["in_primary"] = srcs_df[["primary", "img_list"]].apply(
-        check_primary_image, axis=1
+    # Compute missing images. per source (img_diff_series) and per source detection image
+    img_diff_series, per_source_df = _compute_missing_images(
+        src_skyrg_df, lists_df, img_mult, epoch_mult
     )
-
-    srcs_df = srcs_df.drop(["img_list", "skyreg_img_list", "primary"], axis=1)
+    # img_diff_series: per-source list of missing ideal images.
+    # per_source_df: has per-source 'detection' (first observed image) and 
+    # 'in_primary' (whether the source was detected in its first ideal image) columns.
+    del src_skyrg_df, lists_df
+
+    # Join coords with detection/in_primary, then inner-join img_diff — the
+    # inner join filters down to only sources with >=1 missing image.
+    srcs_df = coords_df.join(per_source_df, how="inner")
+    del coords_df, per_source_df
+    # Join img_diff — the inner join filters down to only sources with >=1 missing image.
+    srcs_df = srcs_df.join(img_diff_series, how="inner")
+    del img_diff_series
+
+    # Convert int32 image codes back to the original image name strings for the final output.
+    srcs_df["detection"] = image_names[srcs_df["detection"].to_numpy()]
+    srcs_df.index = source_ids[srcs_df.index.to_numpy()]
+    srcs_df.index.name = "source"
+
+    # Convert srcs_df to a Dask DataFrame and explode 'img_diff' to one row per missing image.
+    srcs_df = _explode_missing_images_to_dask(srcs_df, image_names, 1)
 
     logger.info("Ideal source coverage time: %.2f seconds", merged_timer.reset())
 

From 355abd0cd6de2d041e3a1af1c9e0b6b9a2d1bca0 Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 13 Aug 2026 17:02:23 +1000
Subject: [PATCH 05/11] Get tests to pass.

---
 vast_pipeline/pipeline/utils.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index ab8477851..606f0eda6 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -29,7 +29,7 @@
 from vast_pipeline.image.utils import open_fits
 from vast_pipeline.utils.utils import (
     eq_to_cart, StopWatch, optimise_numeric, copy_file_or_dir,
-    delete_file_or_dir, generate_shortuuid, UUID_LEN_SOURCE
+    delete_file_or_dir, generate_shortuuid, UUID_LEN_SOURCE, calculate_n_partitions
 )
 from vast_pipeline.models import (
     Band, Image, Run, SkyRegion
@@ -767,8 +767,7 @@ def parallel_groupby_coord(df: dd.DataFrame,) -> Tuple[pd.DataFrame, pd.DataFram
     Produces two separate per-source DataFrames in a single Dask compute pass:
 
     * **coords_df** — lightweight numeric frame (one float per column) used
-      for the AstroPy sky crossmatch: ``wavg_ra``, ``wavg_dec``,
-      ``flux_peak``.
+      for the AstroPy sky crossmatch: ``wavg_ra``, ``wavg_dec``.
     * **lists_df** — heavyweight frame holding Python list columns
       ``img_list`` and ``epoch_list``, only needed for the "missing image"
       computation in ``get_src_skyregion_merged_df``.  Keeping it separate
@@ -1142,9 +1141,8 @@ def _explode_missing_images_to_dask(
 
     Args:
         srcs_df: One row per source, with 'wavg_ra', 'wavg_dec',
-            'flux_peak', 'detection' (image name), 'in_primary' and
-            'img_diff' (a list of int32 image codes), indexed by source
-            ShortUUID.
+            'detection' (image name), 'in_primary' and 'img_diff',
+            indexed by source.
         image_names: Array mapping int32 image code -> original image name.
         n_cpu: Number of available CPUs/workers, used to size the returned
             Dask DataFrame's partitions (via `calculate_n_partitions`). If
@@ -1160,7 +1158,7 @@ def _explode_missing_images_to_dask(
     srcs_df = dd.from_pandas(srcs_df, npartitions=exploded_npartitions)
 
     srcs_df = srcs_df.reset_index()[
-        ["source", "wavg_ra", "wavg_dec", "img_diff", "detection", "in_primary", "flux_peak"]
+        ["source", "wavg_ra", "wavg_dec", "img_diff", "detection", "in_primary"]
     ].explode("img_diff")
 
     def _convert_img_diff_names(partition: pd.DataFrame) -> pd.DataFrame:
@@ -1245,7 +1243,7 @@ def get_src_skyregion_merged_df(
 
 
     coords_df, lists_df = parallel_groupby_coord(sources_df)
-    # coords_df: wavg_ra/wavg_dec/flux_peak, used by _crossmatch_sources_to_skyregions.
+    # coords_df: wavg_ra/wavg_dec, used by _crossmatch_sources_to_skyregions.
     # lists_df: img_list/epoch_list, used by _compute_missing_images.
     del sources_df
 

From cec6cbdd488054d498c19714bf073363db9b0ceb Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Fri, 14 Aug 2026 08:39:46 +1000
Subject: [PATCH 06/11] Avoid assumption of interger epoch

---
 vast_pipeline/pipeline/utils.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index 606f0eda6..73301e4e3 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -883,6 +883,18 @@ def _build_compact_indices(
     # _compute_missing_images' vectorized membership test.
     img_mult = len(image_names)
 
+    # Epoch labels are free-form strings (e.g. '03x'), so they are factorized
+    # into int32 codes on both sides rather than cast to int.
+    epoch_values = pd.unique(images_df["epoch"])
+    epoch_to_idx = pd.Series(np.arange(len(epoch_values), dtype=np.int32), index=epoch_values)
+    images_df = images_df.assign(
+        epoch=epoch_to_idx.reindex(images_df["epoch"]).to_numpy()
+    )
+    sources_df["epoch"] = sources_df["epoch"].map(epoch_to_idx, meta=("epoch", "int32"))
+    # Multiplier for the combined (source, epoch) key used by
+    # _compute_missing_images.
+    epoch_mult = len(epoch_values)
+
     skyreg_img_df = images_df[["skyreg_id", "name", "epoch", "datetime"]].rename(
         columns={
             "name": "skyreg_img_list",
@@ -890,13 +902,6 @@ def _build_compact_indices(
             "datetime": "skyreg_datetime",
         }
     )
-    # int32/int64 downcasts halve these columns' cost across the large
-    # crossmatch expansion in _crossmatch_sources_to_skyregions.
-    skyreg_img_df["skyreg_epoch"] = skyreg_img_df["skyreg_epoch"].astype(np.int32)
-    # Multiplier for the combined (source, epoch) key used by
-    # _compute_missing_images; captured now while skyreg_img_df still holds
-    # the full epoch universe.
-    epoch_mult = int(skyreg_img_df["skyreg_epoch"].max()) + 1
     # Avoids boxing tz-aware Timestamps in the merges/sorts that follow;
     # only used to establish sort order, then dropped.
     skyreg_img_df["skyreg_datetime"] = skyreg_img_df["skyreg_datetime"].astype(np.int64)

From a718e7df4261e19e9a390b476cc924bd782204eb Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Wed, 19 Aug 2026 16:26:09 +1000
Subject: [PATCH 07/11] Do the explode of missing_sources once.

Also tidy up some commetnt etc.
---
 vast_pipeline/pipeline/forced_extraction.py |   3 +-
 vast_pipeline/pipeline/main.py              |  13 +-
 vast_pipeline/pipeline/new_sources.py       |   3 -
 vast_pipeline/pipeline/utils.py             | 174 +++++++-------------
 4 files changed, 65 insertions(+), 128 deletions(-)

diff --git a/vast_pipeline/pipeline/forced_extraction.py b/vast_pipeline/pipeline/forced_extraction.py
index a9f80f9e8..70180a03b 100644
--- a/vast_pipeline/pipeline/forced_extraction.py
+++ b/vast_pipeline/pipeline/forced_extraction.py
@@ -652,8 +652,7 @@ def forced_extraction(
     # | VAST_2118-06A.EPOCH03x.I.fits |  0.165395 | 2019-10-29 10:01:20.500000+00:00 |             319.652 |              -6.2989 |               6.7401 |
     # | VAST_2118-06A.EPOCH02.I.fits  |  0.16323  | 2019-10-30 08:31:20.200000+00:00 |             319.652 |              -6.2989 |               6.7401 |
 
-    # Explode out the img_diff column.
-    extr_df = extr_df.explode("img_diff").reset_index()
+    extr_df = extr_df.reset_index()
     total_to_extract = extr_df.shape[0].compute()
     logger.info("%d forced measurements to extract", total_to_extract)
 
diff --git a/vast_pipeline/pipeline/main.py b/vast_pipeline/pipeline/main.py
index 8014515ff..fcce847f6 100644
--- a/vast_pipeline/pipeline/main.py
+++ b/vast_pipeline/pipeline/main.py
@@ -278,22 +278,11 @@ def process_pipeline(self, p_run: Run) -> None:
             unforced_df,
             images_df,
             skyregs_df,
+            self.dm.num_workers,
         )
         del images_df
         del unforced_df
 
-        # Make missing sources into Dask dataframe
-        # NOTE: This would not be necessary if the get_src_skyregion_merged_df
-        # function was improved to use Dask. (See NOTE in parallel_groupby function.)
-        #npartitions = calculate_n_partitions(missing_sources_df,
-        #                                     n_cpu=self.dm.num_workers,
-        #                                     partition_size_mb=self.config['processing']['max_partition_mb'])
-        #missing_sources_df = dd.from_pandas(
-        #    missing_sources_df,
-        #    npartitions=npartitions
-        #)
-        #wait(missing_sources_df)
-
         # STEP #4 New source analysis
         logger.info("Running step #4: new source analysis...")
         new_sources_df = new_sources(
diff --git a/vast_pipeline/pipeline/new_sources.py b/vast_pipeline/pipeline/new_sources.py
index 40aeccedd..55345a324 100644
--- a/vast_pipeline/pipeline/new_sources.py
+++ b/vast_pipeline/pipeline/new_sources.py
@@ -357,9 +357,6 @@ def new_sources(
     # save the index before exploding
     new_sources_df = new_sources_df.reset_index()
 
-    # Explode now to avoid two loops below
-    new_sources_df = new_sources_df.explode('img_diff')
-
     # Merge the respective image information to the df
     new_sources_df = new_sources_df.merge(
         images_df[['datetime']],
diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index 73301e4e3..2eb97f397 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -850,7 +850,7 @@ def get_rms_noise_image_values(rms_path: str) -> Tuple[float, float, float]:
 
 def _build_compact_indices(
     sources_df: dd.DataFrame, images_df: pd.DataFrame,
-) -> Tuple[dd.DataFrame, np.ndarray, int, pd.DataFrame, int]:
+) -> Tuple[dd.DataFrame, np.ndarray, pd.DataFrame]:
     """
     Replaces image name strings with compact int32 codes for the rest of
     `get_src_skyregion_merged_df`, and builds the per-image ideal-coverage
@@ -859,41 +859,32 @@ def _build_compact_indices(
     Args:
         sources_df: The association step output, must have an 'image'
             column holding image name strings.
-        images_df: All image objects for the run, with 'name', 'skyreg_id',
+        images_df: All images, with 'name', 'skyreg_id',
             'epoch' and 'datetime' columns.
 
     Returns:
         sources_df: With the 'image' column replaced by int32 codes.
         image_names: Array mapping int32 image code -> original image name.
-        img_mult: Multiplier for the combined (source, image) key used by
-            the vectorized "missing image" membership test.
         skyreg_img_df: Per-image ideal-coverage frame indexed by
             'skyreg_id', columns 'skyreg_img_list' (int32 image code),
             'skyreg_epoch' (int32) and 'skyreg_datetime' (int64, a sort key
             only, not a real datetime).
-        epoch_mult: Multiplier for the combined (source, epoch) key.
     """
-    # Compact int32 image code replaces the name string from here on;
+    # int32 image code replaces the name string from here on;
     # image_names converts back to names in get_src_skyregion_merged_df.
     image_names = images_df["name"].to_numpy()
     name_to_idx = pd.Series(np.arange(len(image_names), dtype=np.int32), index=image_names)
     images_df = images_df.assign(name=np.arange(len(image_names), dtype=np.int32))
     sources_df["image"] = sources_df["image"].map(name_to_idx, meta=("image", "int32"))
-    # Multiplier for the combined (source, image) key used by
-    # _compute_missing_images' vectorized membership test.
-    img_mult = len(image_names)
 
-    # Epoch labels are free-form strings (e.g. '03x'), so they are factorized
-    # into int32 codes on both sides rather than cast to int.
+    # int32 epoch code replaces the original epoch string from here on;
+    # epoch_to_idx converts back to original epoch strings in
     epoch_values = pd.unique(images_df["epoch"])
     epoch_to_idx = pd.Series(np.arange(len(epoch_values), dtype=np.int32), index=epoch_values)
     images_df = images_df.assign(
         epoch=epoch_to_idx.reindex(images_df["epoch"]).to_numpy()
     )
     sources_df["epoch"] = sources_df["epoch"].map(epoch_to_idx, meta=("epoch", "int32"))
-    # Multiplier for the combined (source, epoch) key used by
-    # _compute_missing_images.
-    epoch_mult = len(epoch_values)
 
     skyreg_img_df = images_df[["skyreg_id", "name", "epoch", "datetime"]].rename(
         columns={
@@ -902,12 +893,11 @@ def _build_compact_indices(
             "datetime": "skyreg_datetime",
         }
     )
-    # Avoids boxing tz-aware Timestamps in the merges/sorts that follow;
-    # only used to establish sort order, then dropped.
+    # Convert timestamp to lightweight int64;
     skyreg_img_df["skyreg_datetime"] = skyreg_img_df["skyreg_datetime"].astype(np.int64)
     skyreg_img_df = skyreg_img_df.set_index("skyreg_id")
 
-    return sources_df, image_names, img_mult, skyreg_img_df, epoch_mult
+    return sources_df, image_names, skyreg_img_df
 
 
 def _crossmatch_sources_to_skyregions(
@@ -973,8 +963,7 @@ def _crossmatch_sources_to_skyregions(
 def _dedupe_closest_skyregion_per_epoch(src_skyrg_df: pd.DataFrame) -> pd.DataFrame:
     """
     For each (source, ideal epoch) pair, keeps only the closest-matching
-    sky region (smallest separation), then sorts the result ready for the
-    vectorized "missing image" computation that follows.
+    sky region (smallest separation), then sorts the result by datetime then source.
 
     Args:
         src_skyrg_df: Output of `_crossmatch_sources_to_skyregions`, one
@@ -1009,7 +998,7 @@ def _dedupe_closest_skyregion_per_epoch(src_skyrg_df: pd.DataFrame) -> pd.DataFr
     src_skyrg_df.sort_values(by="skyreg_datetime", inplace=True)
     src_skyrg_df.drop(columns=["skyreg_datetime"], inplace=True)
 
-    # Stable sort to preserve the chronological order of each source's rows
+    # Stable sort to preserve the datetime order of each source's rows
     src_skyrg_df.sort_values(by="source", kind="stable", inplace=True)
 
     return src_skyrg_df
@@ -1034,20 +1023,20 @@ def _isin_sorted(keys: np.ndarray, sorted_unique_ref: np.ndarray) -> np.ndarray:
 
 
 def _compute_missing_images(
-    src_skyrg_df: pd.DataFrame, lists_df: pd.DataFrame, img_mult: int, epoch_mult: int,
+    src_skyrg_df: pd.DataFrame, lists_df: pd.DataFrame, img_mult: int, epoch_mult: int
 ) -> Tuple[pd.Series, pd.DataFrame]:
     """
     For each source, determines which ideal images/epochs were never
     actually observed, plus the source's first ideal ("primary") and first
-    observed ("detection") image. Uses vectorized lookups.
+    observed ("detection") image.
 
     Args:
         src_skyrg_df: Output of `_dedupe_closest_skyregion_per_epoch`, one
             row per (source, ideal image), sorted by 'source'.
         lists_df: Per-source 'img_list'/'epoch_list' columns (the actually
             observed images/epochs), indexed by int32 source code.
-        img_mult: Multiplier for the combined (source, image) key.
-        epoch_mult: Multiplier for the combined (source, epoch) key.
+        img_mult: Number of unique images.
+        epoch_mult: Number of unique epochs.
 
     Returns:
         img_diff_series: Per-source list of missing ideal images, indexed
@@ -1059,13 +1048,14 @@ def _compute_missing_images(
     """
     # For each (source, ideal image)/(source, ideal epoch) pair, test
     # whether it was ever actually observed by encoding the pair as one
-    # combined int64 key (source * multiplier + value) and binary-searching
+    # combined key (source * multiplier + value) and binary-searching
     # it (np.searchsorted) against a sorted array of observed keys.
     source_arr = src_skyrg_df["source"].to_numpy()
     skyreg_img_arr = src_skyrg_df["skyreg_img_list"].to_numpy()
     skyreg_epoch_arr = src_skyrg_df["skyreg_epoch"].to_numpy()
     del src_skyrg_df
 
+    # Find the points in source_arr where the source changes
     split_points = np.flatnonzero(np.diff(source_arr)) + 1
     group_start_idx = np.concatenate(([0], split_points))
     group_source = source_arr[group_start_idx]
@@ -1077,6 +1067,7 @@ def _compute_missing_images(
     # exploding the small (one row per source) img_list/epoch_list columns.
     obs_img_long = lists_df["img_list"].explode()
     obs_epoch_long = lists_df["epoch_list"].explode()
+    # Encode each observed pair as a unique int64 key
     obs_img_keys = np.unique(
         obs_img_long.index.to_numpy().astype(np.int64) * img_mult
         + obs_img_long.to_numpy().astype(np.int64)
@@ -1087,14 +1078,16 @@ def _compute_missing_images(
     )
     del obs_img_long, obs_epoch_long
 
+    # Now do the same for the ideal pairs, and test whether each is present in the
+    # observed keys. The result is a boolean array, True where the ideal pair was observed.
     ideal_img_key = source_arr.astype(np.int64) * img_mult + skyreg_img_arr.astype(np.int64)
     in_img_list = _isin_sorted(ideal_img_key, obs_img_keys)
     del ideal_img_key
-
     ideal_epoch_key = source_arr.astype(np.int64) * epoch_mult + skyreg_epoch_arr.astype(np.int64)
     in_epoch_list = _isin_sorted(ideal_epoch_key, obs_epoch_keys)
     del ideal_epoch_key, skyreg_epoch_arr
-
+    # An ideal image only counts as missing if the source was seen neither in that image nor
+    # anywhere in that epoch (avoids false positives from overlapping sky-regions sharing an epoch)
     missing_mask = ~in_img_list & ~in_epoch_list
     del in_img_list, in_epoch_list
 
@@ -1102,7 +1095,7 @@ def _compute_missing_images(
     missing_img = skyreg_img_arr[missing_mask]
     del skyreg_img_arr, missing_mask
 
-    # missing_source is a subset of the already (stable-)sorted source_arr,
+    # missing_source is a subset of the already sorted source_arr,
     # so it's still sorted — group it straight back into per-source lists
     # without needing to re-sort.
     if len(missing_source) > 0:
@@ -1114,6 +1107,8 @@ def _compute_missing_images(
         img_diff_groups = []
     del missing_source, missing_img, source_arr
 
+    # img_diff_series is indexed by source (only sources with ≥1 missing image),
+    # values are arrays of missing images.
     img_diff_series = pd.Series(
         img_diff_groups, index=pd.Index(m_group_source, name="source"), name="img_diff",
     )
@@ -1135,51 +1130,8 @@ def _compute_missing_images(
 
     return img_diff_series, per_source_df
 
-
-def _explode_missing_images_to_dask(
-    srcs_df: pd.DataFrame, image_names: np.ndarray, n_cpu: Optional[int],
-) -> dd.DataFrame:
-    """
-    Converts the assembled per-source result to a Dask DataFrame and
-    explodes 'img_diff' to one row per missing image, pre-computing the
-    exploded form shared by the pipeline's steps #4 and #5.
-
-    Args:
-        srcs_df: One row per source, with 'wavg_ra', 'wavg_dec',
-            'detection' (image name), 'in_primary' and 'img_diff',
-            indexed by source.
-        image_names: Array mapping int32 image code -> original image name.
-        n_cpu: Number of available CPUs/workers, used to size the returned
-            Dask DataFrame's partitions (via `calculate_n_partitions`). If
-            None, partitions are sized on memory alone.
-
-    Returns:
-        Dask DataFrame with one row per (source, missing image) pair, as
-        described in `get_src_skyregion_merged_df`.
-    """
-    base_npartitions = calculate_n_partitions(srcs_df, n_cpu=n_cpu, partition_size_mb=10)
-    exploded_npartitions = max(n_cpu, base_npartitions) if n_cpu else base_npartitions
-
-    srcs_df = dd.from_pandas(srcs_df, npartitions=exploded_npartitions)
-
-    srcs_df = srcs_df.reset_index()[
-        ["source", "wavg_ra", "wavg_dec", "img_diff", "detection", "in_primary"]
-    ].explode("img_diff")
-
-    def _convert_img_diff_names(partition: pd.DataFrame) -> pd.DataFrame:
-        """Vectorized int32-code -> image-name lookup for one partition"""
-        partition = partition.copy()
-        partition["img_diff"] = image_names[partition["img_diff"].to_numpy().astype(np.int32)]
-        return partition
-
-    srcs_df = srcs_df.map_partitions(
-        _convert_img_diff_names, meta=srcs_df._meta.assign(img_diff=pd.Series(dtype=object)),
-    )
-
-    return srcs_df
-
 def get_src_skyregion_merged_df(
-    sources_df: dd.DataFrame, images_df: pd.DataFrame, skyreg_df: pd.DataFrame
+    sources_df: dd.DataFrame, images_df: pd.DataFrame, skyreg_df: pd.DataFrame, n_cpu: int
 ) -> pd.DataFrame:
     """
     Analyses the current sources_df to determine what the 'ideal coverage'
@@ -1196,53 +1148,37 @@ def get_src_skyregion_merged_df(
         skyreg_df:
             Contains the sky regions of the pipeline run. I.e. all
             sky region objects for the run loaded into a dataframe.
+        n_cpu:
+            The number of dask workers.
 
     Returns:
         DataFrame containing missing image information (see source code for
             dataframe format).
     """
     # Output format:
-    # +----------+----------------------------------+-----------+------------+
-    # |   source | img_list                         |   wavg_ra |   wavg_dec |
-    # |----------+----------------------------------+-----------+------------+
-    # |      278 | ['VAST_0127-73A.EPOCH01.I.fits'] |  22.2929  |   -71.8717 |
-    # |      702 | ['VAST_0127-73A.EPOCH01.I.fits'] |  28.8125  |   -69.3547 |
-    # |      844 | ['VAST_0127-73A.EPOCH01.I.fits'] |  17.3152  |   -72.346  |
-    # |      934 | ['VAST_0127-73A.EPOCH01.I.fits'] |   9.75754 |   -72.9629 |
-    # |     1290 | ['VAST_0127-73A.EPOCH01.I.fits'] |  20.8455  |   -76.8269 |
-    # +----------+----------------------------------+-----------+------------+
-    # ------------------------------------------------------------------+
-    #  skyreg_img_list                                                  |
-    # ------------------------------------------------------------------+
-    #  ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
-    #  ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
-    #  ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
-    #  ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
-    #  ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
-    # ------------------------------------------------------------------+
-    # ----------------------------------+------------------------------+
-    #  img_diff                         | primary                      |
-    # ----------------------------------+------------------------------+
-    #  ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
-    #  ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
-    #  ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
-    #  ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
-    #  ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
-    # ----------------------------------+------------------------------+
-    # ------------------------------+--------------+
-    #  detection                    | in_primary   |
-    # ------------------------------+--------------|
-    #  VAST_0127-73A.EPOCH01.I.fits | True         |
-    #  VAST_0127-73A.EPOCH01.I.fits | True         |
-    #  VAST_0127-73A.EPOCH01.I.fits | True         |
-    #  VAST_0127-73A.EPOCH01.I.fits | True         |
-    #  VAST_0127-73A.EPOCH01.I.fits | True         |
-    # ------------------------------+--------------+
+    # +--------------+-----------+------------+
+    # |   source     |   wavg_ra |   wavg_dec |
+    # |--------------+-----------+------------+
+    # | 222BSmZCCu4e |  22.2929  |   -71.8717 |
+    # | 222juojpk42L |  28.8125  |   -69.3547 |
+    # | 222BSmZCCu4f |  17.3152  |   -72.346  |
+    # | 222oo8Gj53r6 |  9.75754  |   -72.9629 |
+    # | g2KSTpuVJthm |  20.8455  |   -76.8269 |
+    # +--------------+-----------+------------+
+    # --------------------------------+------------------------------+--------------+
+    #  img_diff                       | detection                    | in_primary   |
+    # --------------------------------+------------------------------+--------------|
+    #  'VAST_0127-73A.EPOCH08.I.fits' | VAST_0127-73A.EPOCH01.I.fits | True         |
+    #  'VAST_0127-73A.EPOCH08.I.fits' | VAST_0127-73A.EPOCH01.I.fits | True         |
+    #  'VAST_0127-73A.EPOCH08.I.fits' | VAST_0127-73A.EPOCH01.I.fits | True         |
+    #  'VAST_0127-73A.EPOCH08.I.fits' | VAST_0127-73A.EPOCH01.I.fits | True         |
+    #  'VAST_0127-73A.EPOCH08.I.fits' | VAST_0127-73A.EPOCH01.I.fits | True         |
+    # --------------------------------+------------------------------+--------------+
     logger.info("Creating ideal source coverage df...")
 
     merged_timer = StopWatch()
 
-    sources_df, image_names, img_mult, skyreg_img_df, epoch_mult = _build_compact_indices(
+    sources_df, image_names, skyreg_img_df = _build_compact_indices(
         sources_df, images_df
     )
 
@@ -1270,7 +1206,7 @@ def get_src_skyregion_merged_df(
 
     # Compute missing images. per source (img_diff_series) and per source detection image
     img_diff_series, per_source_df = _compute_missing_images(
-        src_skyrg_df, lists_df, img_mult, epoch_mult
+        src_skyrg_df, lists_df, len(image_names), len(pd.unique(images_df["epoch"]))
     )
     # img_diff_series: per-source list of missing ideal images.
     # per_source_df: has per-source 'detection' (first observed image) and 
@@ -1291,7 +1227,23 @@ def get_src_skyregion_merged_df(
     srcs_df.index.name = "source"
 
     # Convert srcs_df to a Dask DataFrame and explode 'img_diff' to one row per missing image.
-    srcs_df = _explode_missing_images_to_dask(srcs_df, image_names, 1)
+    exploded_npartitions = calculate_n_partitions(srcs_df, n_cpu=n_cpu)
+
+    srcs_df = dd.from_pandas(srcs_df, npartitions=exploded_npartitions)
+
+    srcs_df = srcs_df.reset_index()[
+        ["source", "wavg_ra", "wavg_dec", "img_diff", "detection", "in_primary"]
+    ].explode("img_diff")
+
+    def _convert_img_diff_names(partition: pd.DataFrame) -> pd.DataFrame:
+        """Vectorized int32-code -> image-name lookup for one partition"""
+        partition = partition.copy()
+        partition["img_diff"] = image_names[partition["img_diff"].to_numpy().astype(np.int32)]
+        return partition
+
+    srcs_df = srcs_df.map_partitions(
+        _convert_img_diff_names, meta=srcs_df._meta.assign(img_diff=pd.Series(dtype=object)),
+    )
 
     logger.info("Ideal source coverage time: %.2f seconds", merged_timer.reset())
 

From 10a8cff1c9181ce2b695e8f2454be352eea48d0d Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 20 Aug 2026 08:59:12 +1000
Subject: [PATCH 08/11] Tidy up som emore comments.

---
 vast_pipeline/pipeline/utils.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index 2eb97f397..31e9f3cf7 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -1006,7 +1006,8 @@ def _dedupe_closest_skyregion_per_epoch(src_skyrg_df: pd.DataFrame) -> pd.DataFr
 
 def _isin_sorted(keys: np.ndarray, sorted_unique_ref: np.ndarray) -> np.ndarray:
     """
-    Vectorized membership test of `keys` against a sorted, unique reference
+    Optimised implementation of `np.isin(keys, sorted_unique_ref)`. Does a
+    vectorized membership test of `keys` against a sorted, unique reference
     array, via binary search.
 
     Args:
@@ -1182,14 +1183,13 @@ def get_src_skyregion_merged_df(
         sources_df, images_df
     )
 
-
     coords_df, lists_df = parallel_groupby_coord(sources_df)
     # coords_df: wavg_ra/wavg_dec, used by _crossmatch_sources_to_skyregions.
     # lists_df: img_list/epoch_list, used by _compute_missing_images.
     del sources_df
 
     # Use compact int32 source code instead of the ShortUUID string for the
-    # rest of this function; converted back on the final index only.
+    # rest of this function.
     source_ids = coords_df.index.to_numpy()
     source_to_idx = pd.Series(np.arange(len(source_ids), dtype=np.int32), index=source_ids)
     coords_df.index = np.arange(len(source_ids), dtype=np.int32)
@@ -1217,7 +1217,6 @@ def get_src_skyregion_merged_df(
     # inner join filters down to only sources with >=1 missing image.
     srcs_df = coords_df.join(per_source_df, how="inner")
     del coords_df, per_source_df
-    # Join img_diff — the inner join filters down to only sources with >=1 missing image.
     srcs_df = srcs_df.join(img_diff_series, how="inner")
     del img_diff_series
 
@@ -1228,9 +1227,7 @@ def get_src_skyregion_merged_df(
 
     # Convert srcs_df to a Dask DataFrame and explode 'img_diff' to one row per missing image.
     exploded_npartitions = calculate_n_partitions(srcs_df, n_cpu=n_cpu)
-
     srcs_df = dd.from_pandas(srcs_df, npartitions=exploded_npartitions)
-
     srcs_df = srcs_df.reset_index()[
         ["source", "wavg_ra", "wavg_dec", "img_diff", "detection", "in_primary"]
     ].explode("img_diff")

From f6673952d2f824ddfa6620ce411225d34f603cda Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 20 Aug 2026 09:37:38 +1000
Subject: [PATCH 09/11] Persist output dataframe and fix type hints

---
 vast_pipeline/pipeline/utils.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index 31e9f3cf7..c43040c7b 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -18,6 +18,8 @@
 import tempfile
 import itertools
 
+from dask.distributed import wait
+
 from typing import Any, List, Optional, Dict, Tuple, Union
 from astropy.coordinates import SkyCoord, Angle
 from django.conf import settings
@@ -1133,7 +1135,7 @@ def _compute_missing_images(
 
 def get_src_skyregion_merged_df(
     sources_df: dd.DataFrame, images_df: pd.DataFrame, skyreg_df: pd.DataFrame, n_cpu: int
-) -> pd.DataFrame:
+) -> dd.DataFrame:
     """
     Analyses the current sources_df to determine what the 'ideal coverage'
     for each source should be. In other words, what images is the source
@@ -1242,6 +1244,9 @@ def _convert_img_diff_names(partition: pd.DataFrame) -> pd.DataFrame:
         _convert_img_diff_names, meta=srcs_df._meta.assign(img_diff=pd.Series(dtype=object)),
     )
 
+    srcs_df = srcs_df.persist()
+    wait(srcs_df)
+
     logger.info("Ideal source coverage time: %.2f seconds", merged_timer.reset())
 
     return srcs_df

From 4e82fe9a4bfea9eaa0ca9cef385e0fcfe13af62e Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 20 Aug 2026 09:46:41 +1000
Subject: [PATCH 10/11] Linting.

---
 vast_pipeline/pipeline/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/vast_pipeline/pipeline/utils.py b/vast_pipeline/pipeline/utils.py
index c43040c7b..decb071e7 100644
--- a/vast_pipeline/pipeline/utils.py
+++ b/vast_pipeline/pipeline/utils.py
@@ -20,7 +20,7 @@
 
 from dask.distributed import wait
 
-from typing import Any, List, Optional, Dict, Tuple, Union
+from typing import Any, List, Optional, Dict, Tuple
 from astropy.coordinates import SkyCoord, Angle
 from django.conf import settings
 from django.contrib.auth.models import User
@@ -1211,7 +1211,7 @@ def get_src_skyregion_merged_df(
         src_skyrg_df, lists_df, len(image_names), len(pd.unique(images_df["epoch"]))
     )
     # img_diff_series: per-source list of missing ideal images.
-    # per_source_df: has per-source 'detection' (first observed image) and 
+    # per_source_df: has per-source 'detection' (first observed image) and
     # 'in_primary' (whether the source was detected in its first ideal image) columns.
     del src_skyrg_df, lists_df
 

From e6208d5f253314eb93c86f7ae2302f49264fcf3c Mon Sep 17 00:00:00 2001
From: Tom Mauch <thomas.mauch@sydney.edu.au>
Date: Thu, 20 Aug 2026 09:49:32 +1000
Subject: [PATCH 11/11] Add changelog entry.

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ddab4e255..cb7e49015 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 #### Changed
 
+- V2: Optimise missing source calculation to use vectorised numpy operations [#899](https://github.com/askap-vast/vast-pipeline/pull/899)
 - V2: Pin package repos for compatibility with latest Python 3.11 & 3.12 [#879](https://github.com/askap-vast/vast-pipeline/pull/879)
 - V2: Revert back to parallel association upload [#852](https://github.com/askap-vast/vast-pipeline/pull/852)
 - V2: Change worker throttling to use Dask Semaphore [#852](https://github.com/askap-vast/vast-pipeline/pull/852)
@@ -95,6 +96,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Removed Aladin lite display containing RACS and other ASKAP HIPS images [#874](https://github.com/askap-vast/vast-pipeline/pull/874)
 
 #### List of PRs
+- [#899](https://github.com/askap-vast/vast-pipeline/pull/899): feat: V2: Optimise missing source calculation
 - [#893](https://github.com/askap-vast/vast-pipeline/pull/893): feat: V2: Parallelise ingest of measurements
 - [#880](https://github.com/askap-vast/vast-pipeline/pull/880): fix: V2: Ensure external dependencies pass when etxternal websites are down
 - [#879](https://github.com/askap-vast/vast-pipeline/pull/879): fix: V2: Update dependencies to work with latest python 3.11&3.12 and pin them