From 5de5f10ce2f008396e7374263eb40d6336266e22 Mon Sep 17 00:00:00 2001
From: Brian Healy <42810347+bfhealy@users.noreply.github.com>
Date: Mon, 8 Apr 2024 19:36:39 -0500
Subject: [PATCH] Add check-quads-for-sources script from
 generate_features_slurm (#575)

---
 pyproject.toml                   |  1 +
 tools/generate_features_slurm.py | 51 ++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 8adb4ef..db4c28d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -79,6 +79,7 @@ scope-manage-annotation = "tools.scope_manage_annotation:main"
 post-taxonomy = "tools.taxonomy:main"
 generate-features = "tools.generate_features:main"
 generate-features-slurm = "tools.generate_features_slurm:main"
+check-quads-for-sources = "tools.generate_features_slurm:check_quads_for_sources"
 generate-features-job-submission = "tools.generate_features_job_submission:main"
 train-algorithm-slurm = "tools.train_algorithm_slurm:main"
 train-algorithm-job-submission = "tools.train_algorithm_job_submission:main"
diff --git a/tools/generate_features_slurm.py b/tools/generate_features_slurm.py
index 24f9c7d..26c3775 100755
--- a/tools/generate_features_slurm.py
+++ b/tools/generate_features_slurm.py
@@ -78,6 +78,16 @@ def check_quads_for_sources(
     :return missing_ccd_quad: boolean stating whether each field in fields has no sources in at least one ccd/quad
     """
 
+    parser = get_check_quads_parser()
+    args, _ = parser.parse_known_args()
+
+    fields = args.fields
+    catalog = args.catalog
+    count_sources = args.count_sources
+    minobs = args.minobs
+    save = args.save
+    filename = args.filename
+
     running_total_sources = 0
     has_sources = np.zeros(len(fields), dtype=bool)
     missing_ccd_quad = np.zeros(len(fields), dtype=bool)
@@ -176,6 +186,47 @@ def check_quads_for_sources(
     return field_dct, has_sources, missing_ccd_quad
 
 
+def get_check_quads_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--fields",
+        type=int,
+        nargs="+",
+        default=np.arange(0, 2000),
+        help="list of integer field numbers to query",
+    )
+    parser.add_argument(
+        "--catalog",
+        type=str,
+        default=source_catalog,
+        help="name of source catalog to query",
+    )
+    parser.add_argument(
+        "--count-sources",
+        action="store_true",
+        help="if set, count number of sources per quad and return",
+    )
+    parser.add_argument(
+        "--minobs",
+        type=int,
+        default=0,
+        help="minimum number of observations needed to count a source",
+    )
+    parser.add_argument(
+        "--save",
+        action="store_true",
+        help="if set, save results dictionary in json format",
+    )
+    parser.add_argument(
+        "--filename",
+        type=str,
+        default="catalog_completeness",
+        help="filename of saved results",
+    )
+
+    return parser
+
+
 def get_slurm_parser():
 
     fg_parser = get_parser(add_help=False)