Refactoring

Author: erinleighh

bin.src/curate_templates.py

@@ -29,11 +29,10 @@
 
 import argparse
 import logging
-import numpy as np
 import os
 import sys
 
-from astropy.table import Table
+from astropy.table import Table, vstack
 
 from lsst.daf.butler import Butler, CollectionType
 
@@ -56,6 +55,7 @@ def _make_parser():
         "--collections",
         action="extend",
         nargs="+",
+        required=True,
         help="The input collections to search for template_coadd and coadd_depth_table datasets.",
     )
     parser.add_argument(
@@ -86,20 +86,20 @@ def _make_parser():
         "--cutoff",
         required=False,
         default=95,
+        type=int,
         help="The curation process will filter out anything below this cutoff."
         " Default is 95.",
     )
     return parser
 
 
-def get_tracts(butler, where):
-    tracts = []
-    coadd_depth_tables = butler.registry.queryDatasets(datasetType='coadd_depth_table', where=where)
-    for item in coadd_depth_tables:
-        tract = item.dataId['tract']
-        tracts.append(tract)
-    tracts = set(tracts)
-    return tracts
+def select_ref(drefs, tract, patch, band, dtype="template_coadd"):
+    if not drefs:
+        logging.warning(f"No {dtype} found for tract {tract}, patch {patch}, band {band}. Skipping.")
+        return None
+    if len(drefs) > 1:
+        return sorted(drefs, key=lambda ref: ref.run)[-1]
+    return drefs[0]
 
 
 def make_threshold_cuts(butler, template_coadds, n_images, tracts, filter_by, cutoff):
@@ -126,13 +126,21 @@ def make_threshold_cuts(butler, template_coadds, n_images, tracts, filter_by, cu
                             and d.dataId['band'] == band
                             ]
 
+            if not dref:
+                logging.warning(f"No template_coadd found for tract {tract}, patch {patch}, band {band}. "
+                                f"Skipping.")
+                continue
             if len(dref) > 1:
                 sorted_dupe_entry = sorted(dref, key=lambda ref: ref.run)
                 ref = sorted_dupe_entry[-1]
             else:
                 ref = dref[0]
             accepted_drefs.append(ref)
 
+            if not n_image_dref:
+                logging.warning(f"No template_coadd_n_image found for tract {tract}, patch {patch}, "
+                                f"band {band}. Skipping.")
+                continue
             if len(n_image_dref) > 1:
                 sorted_dupe_entry = sorted(n_image_dref, key=lambda ref: ref.run)
                 n_image_ref = sorted_dupe_entry[-1]
@@ -150,6 +158,10 @@ def make_threshold_cuts(butler, template_coadds, n_images, tracts, filter_by, cu
                     and d.dataId['band'] == band
                     ]
 
+            if not dref:
+                logging.warning(f"No template_coadd found for tract {tract}, patch {patch}, band {band}. "
+                                f"Skipping.")
+                continue
             if len(dref) > 1:
                 sorted_dupe_entry = sorted(dref, key=lambda ref: ref.run)
                 ref = sorted_dupe_entry[-1]
@@ -160,68 +172,75 @@ def make_threshold_cuts(butler, template_coadds, n_images, tracts, filter_by, cu
 
 
 def run_stats(accepted_drefs, rejected_drefs, tracts, stats_records_file):
-    # Create table of accepted drefs
-    accepted = Table()
-    accepted_tracts = []
-    accepted_patches = []
-    accepted_bands = []
-    bands = ['u', 'g', 'r', 'i', 'z', 'y']
-
-    for ref in accepted_drefs:
-        accepted_tracts.append(ref.dataId['tract'])
-        accepted_patches.append(ref.dataId['patch'])
-        accepted_bands.append(ref.dataId['band'])
-
-    accepted_table_data = [accepted_tracts, accepted_patches, accepted_bands]
-    accepted = Table(data=accepted_table_data, names=['tract', 'patch', 'band'])
+    """
+    Compute per-tract and per-band accepted/rejected statistics and save to CSV.
+
+    Parameters
+    ----------
+    accepted_drefs : list of DatasetRef
+        Template coadd references that passed curation.
+    rejected_drefs : list of DatasetRef
+        Template coadd references that failed curation.
+    tracts : iterable of int
+        List of tract IDs to include in the stats.
+    stats_records_file : str
+        Path to save the resulting CSV file.
+    """
 
-    # Create table of rejected drefs
-    rejected = Table()
-    rejected_tracts = []
-    rejected_patches = []
-    rejected_bands = []
-
-    for ref in rejected_drefs:
-        rejected_tracts.append(ref.dataId['tract'])
-        rejected_patches.append(ref.dataId['patch'])
-        rejected_bands.append(ref.dataId['band'])
+    bands = ['u', 'g', 'r', 'i', 'z', 'y']
 
-    rejected_table_data = [rejected_tracts, rejected_patches, rejected_bands]
-    rejected = Table(data=rejected_table_data, names=['tract', 'patch', 'band'])
+    # Build accepted table
+    if accepted_drefs:
+        accepted = Table(
+            {
+                'tract': [int(r.dataId['tract']) for r in accepted_drefs],
+                'patch': [int(r.dataId['patch']) for r in accepted_drefs],
+                'band': [str(r.dataId['band']) for r in accepted_drefs],
+                'status': ['accepted'] * len(accepted_drefs)
+            }
+        )
+    else:
+        accepted = Table(names=('tract', 'patch', 'band', 'status'))
+
+    # Build rejected table
+    if rejected_drefs:
+        rejected = Table(
+            {
+                'tract': [int(r.dataId['tract']) for r in rejected_drefs],
+                'patch': [int(r.dataId['patch']) for r in rejected_drefs],
+                'band': [str(r.dataId['band']) for r in rejected_drefs],
+                'status': ['rejected'] * len(rejected_drefs)
+            }
+        )
+    else:
+        rejected = Table(names=('tract', 'patch', 'band', 'status'))
+
+    # Combine tables
+    all_refs = vstack([accepted, rejected])
+
+    # Group by tract and band
+    grouped = all_refs.group_by(['tract', 'band'])
+
+    # Prepare output table
+    stat_table_data = {'tract': [], }
+    for band in bands:
+        stat_table_data[f'{band}_num_accepted'] = []
+        stat_table_data[f'{band}_percent_accepted'] = []
 
-    # Run stats
-    by_band_stats = []
     for tract in tracts:
-        tract_band_stats = []
-        for band in bands:
-            accepted_bands = ((accepted['tract'] == tract) & (accepted['band'] == band)).sum()
-            rejected_bands = ((rejected['tract'] == tract) & (rejected['band'] == band)).sum()
-            total_bands = accepted_bands + rejected_bands
-            if total_bands == 0:
-                tract_band_stats.append(["0 / 0", np.nan])
-            else:
-                tract_band_stats.append([f"{accepted_bands} / {total_bands}",
-                                         accepted_bands / total_bands * 100])
-        by_band_stats.append(tract_band_stats)
-
-    # Compile stats into a table and save
-    accepted_col_names = [f"{band}_{suffix}" for band in bands for suffix
-                          in ("num_accepted", "percent_accepted")]
-    by_tract_names = ['tract'] + accepted_col_names
-
-    stat_table_data = {col: [] for col in by_tract_names}
-
-    for tract_index, tract in enumerate(tracts):
-        band_stats = by_band_stats[tract_index]
-
         stat_table_data['tract'].append(tract)
+        for band in bands:
+            mask = (grouped['tract'] == tract) & (grouped['band'] == band)
+            subset = grouped[mask]
+            n_total = len(subset)
+            n_accepted = (subset['status'] == 'accepted').sum() if n_total > 0 else 0
+            percent = (n_accepted / n_total * 100) if n_total > 0 else float('nan')
+            stat_table_data[f'{band}_num_accepted'].append(f"{n_accepted} / {n_total}")
+            stat_table_data[f'{band}_percent_accepted'].append(percent)
 
-        for band_idx, band in enumerate(bands):
-            accepted_str, percent = band_stats[band_idx]
-            stat_table_data[f"{band}_num_accepted"].append(accepted_str)
-            stat_table_data[f"{band}_percent_accepted"].append(percent)
-    by_tract_stats = Table(stat_table_data)
-    by_tract_stats.write(stats_records_file, format='ascii.csv', overwrite=True)
+    # Create final table
+    stat_table = Table(stat_table_data)
+    stat_table.write(stats_records_file, format='ascii.csv', overwrite=True)
 
 
 def main():
@@ -254,21 +273,39 @@ def main():
         logging.error(f"Collection {tagged_collection} already exists. Aborting.")
         sys.exit(1)
 
-    logging.info("Collecting template_coadd and template_coadd_n_image refs.")
-    refs = butler.query_datasets("template_coadd", where=args.where, limit=None)
-    n_image_refs = butler.query_datasets("template_coadd_n_image", where=args.where, limit=None)
-    logging.info(f"Found {len(refs)} template_coadd datasets in {args.collections}.")
+    logging.info("Collecting coadd_depth_table, template_coadd, and template_coadd_n_image refs.")
+    coadd_depth_table_refs = butler.query_datasets("coadd_depth_table", where=args.where, limit=None)
+    if not coadd_depth_table_refs:
+        logging.error("No coadd_depth_table datasets found in the given collections.")
+        sys.exit(1)
+
+    # Get a list of relavent tracts.
+    tracts = {item.dataId['tract'] for item in coadd_depth_table_refs}
+
+    # Ammend the where argument to restrict refs to relavent tracts.
+    tracts_str = ",".join(str(t) for t in tracts)
+    tract_restriction = f"tract IN ({tracts_str})"
+    args.where = f"({args.where}) AND ({tract_restriction})" if args.where else tract_restriction
 
-    # Get a list of the tracts inside the template collection.
-    tracts = get_tracts(butler, args.where)
+    # Get relavent template_coadd and template_coadd_n_image refs.
+    coadd_refs = butler.query_datasets("template_coadd", where=args.where, limit=None)
+    if not coadd_refs:
+        logging.error("No template_coadd datasets found in the given collections.")
+        sys.exit(1)
+    n_image_refs = butler.query_datasets("template_coadd_n_image", where=args.where, limit=None)
+    if not n_image_refs:
+        logging.error("No template_coadd_n_image datasets found in the given collections.")
+        sys.exit(1)
+    logging.info(f"Found {len(coadd_refs)} template_coadd datasets with coadd_depth_tables "
+                 f"in {args.collections}.")
 
     # Filter out template_coads that don't meet the cutoff and save them to record.
     logging.info("Starting curation.")
-    accepted_drefs, rejected_drefs, accepted_n_image_refs = make_threshold_cuts(butler, refs,
+    accepted_drefs, rejected_drefs, accepted_n_image_refs = make_threshold_cuts(butler, coadd_refs,
                                                                                 n_image_refs, tracts,
                                                                                 args.filter_by, args.cutoff
                                                                                 )
-    logging.info(f"Curation complete. Accepted {len(accepted_drefs)} out of {len(refs)}"
+    logging.info(f"Curation complete. Accepted {len(accepted_drefs)} out of {len(coadd_refs)}"
                  f" template_coadd datasets in {args.collections}.")
 
     # Run accepted/rejected statistics and save them to record.
@@ -278,13 +315,10 @@ def main():
     logging.info("Stat generation complete. Accepted/rejected stat records written to"
                  f" {stats_records_file}.")
 
-    # Associate accepted template_coadds to tagged collection.
-    logging.info(f"Associating {len(accepted_drefs)} template_coadds to {tagged_collection}.")
+    # Associate accepted template_coadds and template_coadd_n_images to tagged collection.
+    logging.info(f"Associating {len(accepted_drefs)} template_coadds and "
+                 f"{len(accepted_n_image_refs)} template_coadd_n_images to {tagged_collection}.")
     butler_write.registry.associate(tagged_collection, accepted_drefs)
-    logging.info("Association complete.")
-
-    # Associate accepted template_coadd_n_images to tagged collection.
-    logging.info(f"Associating {len(accepted_n_image_refs)} template_coadd_n_images to {tagged_collection}.")
     butler_write.registry.associate(tagged_collection, accepted_n_image_refs)
     logging.info("Association complete.")