From 84afd3a0aae8bcb8575123f22bd65e90c0baf8aa Mon Sep 17 00:00:00 2001
From: Teoh Zhixiang <47135468+zhixiangteoh@users.noreply.github.com>
Date: Fri, 18 Aug 2023 14:33:52 -0400
Subject: [PATCH] Prune large DECam CCDs PSF params file (#931)

Fixes #922.
---
 .../ccds-annotated-decam-dr9-small.fits       |  3 ++
 scripts/prune_decam_ccds_annotated.py         | 48 +++++++++++++++++++
 tests/conftest.py                             | 23 +++++++++
 tests/test_decals.py                          |  4 --
 4 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 data/tests/decals/ccds-annotated-decam-dr9-small.fits
 create mode 100644 scripts/prune_decam_ccds_annotated.py
 delete mode 100644 tests/test_decals.py

diff --git a/data/tests/decals/ccds-annotated-decam-dr9-small.fits b/data/tests/decals/ccds-annotated-decam-dr9-small.fits
new file mode 100644
index 000000000..60405c965
--- /dev/null
+++ b/data/tests/decals/ccds-annotated-decam-dr9-small.fits
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06191b98602b975ba0b8704637ca1031b82761dcdec4689c521c96aa6e055963
+size 60851520
diff --git a/scripts/prune_decam_ccds_annotated.py b/scripts/prune_decam_ccds_annotated.py
new file mode 100644
index 000000000..a35ec1fbc
--- /dev/null
+++ b/scripts/prune_decam_ccds_annotated.py
@@ -0,0 +1,48 @@
+import numpy as np
+from astropy.io import fits
+from astropy.table import Table
+
+# Script to reduce the size of the ccds-annotated-decam file by only keeping fixed CCDs used for
+# DECaLS PSF params. Use to regenerate data/tests/decals/ccds-annotated-decam-dr9-small.fits.
+
+with open("/home/zhteoh/bliss/data/decals/ccds-annotated-decam-dr9.fits", "rb") as f:
+    ccds_annotated_fits = fits.open(f)
+
+ccds_annotated_table = Table.read("/home/zhteoh/bliss/data/decals/ccds-annotated-decam-dr9.fits")
+
+BRICKNAME = "3366m010"
+brick_ccds = Table.read(
+    f"/home/zhteoh/bliss/data/decals/{BRICKNAME[:3]}/{BRICKNAME}/legacysurvey-{BRICKNAME}-ccds.fits"
+)
+fixed_ccds = brick_ccds[
+    (
+        brick_ccds["image_filename"]
+        == "decam/CP/V4.8.2a/CP20141020/c4d_141021_015854_ooi_g_ls9.fits.fz"
+    )
+    | (
+        brick_ccds["image_filename"]
+        == "decam/CP/V4.8.2a/CP20151107/c4d_151108_003333_ooi_r_ls9.fits.fz"
+    )
+    | (
+        brick_ccds["image_filename"]
+        == "decam/CP/V4.8.2a/CP20130912/c4d_130913_040652_ooi_z_ls9.fits.fz"
+    )
+]
+keep_ccds = fixed_ccds["ccdname"]
+
+psf_cols = [
+    col
+    for col in ccds_annotated_table.colnames
+    if col.startswith("psf") or col.startswith("gal") or col.startswith("gauss")
+]
+brick_ccds_mask = np.isin(ccds_annotated_table["ccdname"], keep_ccds)
+keep_cols = psf_cols + ["ccdname", "filter"]
+ccds_annotated_table_small = ccds_annotated_table[brick_ccds_mask][keep_cols]
+
+# write ccds_annotated_table_small as astropy table to file
+save_location = "/home/zhteoh/bliss/data/tests/decals/ccds-annotated-decam-dr9-small.fits"
+with open(save_location, "wb") as f:
+    ccds_annotated_table_small.write(f, format="fits")
+
+# du -sh /home/zhteoh/bliss/data/decals/ccds-annotated-decam-dr9.fits => 3.8G
+# du -sh /home/zhteoh/bliss/data/tests/decals/ccds-annotated-decam-dr9-small.fits => 59M
diff --git a/tests/conftest.py b/tests/conftest.py
index 826b2f024..4abf6a1d0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,6 @@
 # pylint: skip-file
+import os
+import shutil
 from pathlib import Path
 
 import pytest
@@ -47,6 +49,27 @@ def cfg(pytestconfig, cached_data_path, output_path):
     return the_cfg
 
 
+@pytest.fixture(scope="session")
+def decals_setup_teardown(cfg):
+    # replace if needed
+    original_ccds_annotated_path = cfg.paths.decals + "/ccds-annotated-decam-dr9.fits"
+    temp_ccds_annotated_path = cfg.paths.decals + "/ccds-annotated-decam-dr9-large.fits"
+    large_file_existed = os.path.exists(original_ccds_annotated_path)
+    if large_file_existed:
+        shutil.move(original_ccds_annotated_path, temp_ccds_annotated_path)
+    shutil.copyfile(
+        cfg.paths.data + "/tests/decals/ccds-annotated-decam-dr9-small.fits",
+        original_ccds_annotated_path,
+    )
+
+    yield
+
+    # restore
+    os.remove(original_ccds_annotated_path)
+    if large_file_existed:
+        shutil.move(temp_ccds_annotated_path, original_ccds_annotated_path)
+
+
 @pytest.fixture(scope="session")
 def encoder(cfg):
     encoder = instantiate(cfg.encoder).to(cfg.predict.device)
diff --git a/tests/test_decals.py b/tests/test_decals.py
deleted file mode 100644
index 5123d478d..000000000
--- a/tests/test_decals.py
+++ /dev/null
@@ -1,4 +0,0 @@
-class TestDecals:
-    def test_decals(self, cfg):
-        # TODO: analogous to TestSDSS::test_sdss, but for decals
-        pass