Add resumable all-sky 2MASS acquisition

scripts/download_2mass_psc_all_sky.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+"""Resumable 2MASS PSC acquisition into fixed one-degree output tiles.
+
+The renderer-facing layout is always ``tile_raRRR_decDDD.csv``. For download
+efficiency, one declination band is covered by wider integer-RA sectors at high
+latitude. Every sector is queried by one <= 1 degree IRSA cone and owns a
+disjoint run of final one-degree tiles, so overlap between cones cannot produce
+duplicate output records. Raw tables and photometry intermediates live under
+``/tmp/2mass_psc_all_sky`` and are removed after a successful cover.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import math
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
+from dataclasses import dataclass
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+OUTPUT_ROOT = ROOT / "assets/2mass/processed/all_sky"
+DONE_ROOT = OUTPUT_ROOT / ".done"
+TMP_ROOT = Path("/tmp/2mass_psc_all_sky")
+PROCESSOR = ROOT / "scripts/process_2mass_psc.py"
+RADIUS_DEG = 1.0
+
+
+@dataclass(frozen=True)
+class Tile:
+    dec_index: int
+    ra_index: int
+
+    @property
+    def name(self) -> str:
+        return f"ra{self.ra_index:03d}_dec{self.dec_index:03d}"
+
+    @property
+    def output_name(self) -> str:
+        return f"tile_ra{self.ra_index:03d}_dec{self.dec_index:03d}.csv"
+
+
+@dataclass(frozen=True)
+class Cover:
+    """A cone and its disjoint, integer-degree RA ownership interval."""
+    dec_index: int
+    cover_index: int
+    ra_start: int
+    ra_end: int
+
+    @property
+    def name(self) -> str:
+        return f"dec{self.dec_index:03d}_cover{self.cover_index:03d}_ra{self.ra_start:03d}_{self.ra_end:03d}"
+
+    @property
+    def center_ra(self) -> float:
+        return (self.ra_start + self.ra_end) * 0.5
+
+    @property
+    def center_dec(self) -> float:
+        return self.dec_index - 89.5
+
+
+def tile_indices(ra: float, dec: float) -> tuple[int, int]:
+    """Return directly-computable (ra_index, dec_index) for ICRS degrees."""
+    if not math.isfinite(ra) or not math.isfinite(dec) or dec < -90.0 or dec > 90.0:
+        raise ValueError("expected finite ICRS RA and Dec with -90 <= Dec <= 90")
+    return int(math.floor(ra % 360.0)), min(179, int(math.floor(dec + 90.0)))
+
+
+def covers_for_band(dec_index: int) -> list[Cover]:
+    """Partition a band into maximum safe integer-degree RA cover sectors."""
+    if not 0 <= dec_index < 180:
+        raise ValueError("dec_index must be in 0..179")
+    dec_min = dec_index - 90.0
+    dec_max = dec_min + 1.0
+    # Start from a tangent-plane bound using the edge farther from the pole,
+    # then prove the exact spherical four-corner condition below.
+    edge_cosine = max(math.cos(math.radians(dec_min)), math.cos(math.radians(dec_max)))
+    width = min(360, max(1, int(math.floor(math.sqrt(3.0) / edge_cosine))))
+
+    def corners_fit(candidate: int) -> bool:
+        center_ra = candidate * 0.5
+        center_dec = math.radians(dec_index - 89.5)
+        for ra in (0.0, float(candidate)):
+            for dec in (dec_min, dec_max):
+                source_dec = math.radians(dec)
+                cosine_distance = (math.sin(center_dec) * math.sin(source_dec) +
+                                   math.cos(center_dec) * math.cos(source_dec) *
+                                   math.cos(math.radians(ra - center_ra)))
+                if math.degrees(math.acos(max(-1.0, min(1.0, cosine_distance)))) > RADIUS_DEG:
+                    return False
+        return True
+
+    while width > 1 and not corners_fit(width):
+        width -= 1
+    return [Cover(dec_index, number, start, min(start + width, 360))
+            for number, start in enumerate(range(0, 360, width))]
+
+
+def done_path(tile: Tile) -> Path:
+    return DONE_ROOT / (tile.name + ".done")
+
+
+def tile_is_done(tile: Tile) -> bool:
+    return done_path(tile).exists()
+
+
+def query_cover(cover: Cover, raw_path: Path, outrows: int) -> None:
+    """Fetch and validate one IPAC table without exposing partial downloads."""
+    raw_path.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(prefix=cover.name + "_", suffix=".tbl",
+                                     dir=raw_path.parent, delete=False) as temporary:
+        temporary_path = Path(temporary.name)
+    command = [
+        "curl", "--fail", "--silent", "--show-error", "--location", "--get",
+        "https://irsa.ipac.caltech.edu/cgi-bin/Gator/nph-query",
+        "--data-urlencode", "catalog=fp_psc",
+        "--data-urlencode", "spatial=cone",
+        "--data-urlencode", f"objstr={cover.center_ra:.10f} {cover.center_dec:.10f}",
+        "--data-urlencode", f"radius={RADIUS_DEG}",
+        "--data-urlencode", "radunits=deg",
+        "--data-urlencode", "outfmt=1",
+        "--data-urlencode", "selcols=designation,ra,dec,j_m,h_m,k_m,rd_flg",
+        "--data-urlencode", f"outrows={outrows}", "-o", str(temporary_path),
+    ]
+    try:
+        subprocess.run(command, check=True)
+        text = temporary_path.read_text(encoding="ascii")
+        header = re.search(r"^\|\s*designation\s*\|", text, re.MULTILINE)
+        if "2MASS All-Sky Point Source Catalog" not in text or header is None:
+            raise RuntimeError(f"{cover.name}: IRSA returned no valid PSC IPAC table")
+        row_count = re.search(r"^\\RowsRetrieved\s*=\s*(\d+)\s*$", text, re.MULTILINE)
+        if row_count is not None and int(row_count.group(1)) >= outrows:
+            raise RuntimeError(f"{cover.name}: reached --outrows={outrows}; rerun with a larger cap")
+        temporary_path.replace(raw_path)
+    except Exception:
+        temporary_path.unlink(missing_ok=True)
+        raise
+
+
+def write_tile(tile: Tile, rows: list[dict[str, str]]) -> None:
+    """Atomically write one complete final tile and only then mark it done."""
+    output_path = OUTPUT_ROOT / tile.output_name
+    with tempfile.NamedTemporaryFile(prefix=tile.name + "_", suffix=".csv",
+                                     dir=OUTPUT_ROOT, mode="w", encoding="ascii",
+                                     newline="", delete=False) as temporary:
+        writer = csv.DictWriter(temporary,
+                                fieldnames=("ra_deg", "dec_deg", "temperature_K", "amplitude_sr"),
+                                lineterminator="\n")
+        writer.writeheader()
+        writer.writerows(rows)
+        temporary_path = Path(temporary.name)
+    temporary_path.replace(output_path)
+    done_path(tile).write_text(f"retained_rows={len(rows)}\n", encoding="ascii")
+
+
+def process_cover(cover: Cover, raw_path: Path, allowed_ra: set[int] | None = None) -> int:
+    """Clean a cone then distribute its owned rows into final 1-degree tiles."""
+    with tempfile.NamedTemporaryFile(prefix=cover.name + "_", suffix=".csv",
+                                     dir=TMP_ROOT, delete=False) as temporary:
+        intermediate = Path(temporary.name)
+    try:
+        subprocess.run([sys.executable, str(PROCESSOR), "--input", str(raw_path),
+                        "--output", str(intermediate)], check=True)
+        buckets: dict[int, list[dict[str, str]]] = {}
+        with intermediate.open(newline="", encoding="ascii") as source:
+            for row in csv.DictReader(source):
+                try:
+                    ra_index, dec_index = tile_indices(float(row["ra_deg"]), float(row["dec_deg"]))
+                except ValueError:
+                    continue
+                if dec_index != cover.dec_index or not cover.ra_start <= ra_index < cover.ra_end:
+                    continue
+                if allowed_ra is not None and ra_index not in allowed_ra:
+                    continue
+                tile = Tile(dec_index, ra_index)
+                if not tile_is_done(tile):
+                    buckets.setdefault(ra_index, []).append(row)
+        completed = 0
+        for ra_index in range(cover.ra_start, cover.ra_end):
+            if allowed_ra is not None and ra_index not in allowed_ra:
+                continue
+            tile = Tile(cover.dec_index, ra_index)
+            if not tile_is_done(tile):
+                write_tile(tile, buckets.get(ra_index, []))
+                completed += 1
+        return completed
+    finally:
+        intermediate.unlink(missing_ok=True)
+
+
+def cover_has_work(cover: Cover, allowed_ra: set[int] | None = None) -> bool:
+    return any(not tile_is_done(Tile(cover.dec_index, ra_index))
+               and (allowed_ra is None or ra_index in allowed_ra)
+               for ra_index in range(cover.ra_start, cover.ra_end))
+
+
+def download_cover(cover: Cover, outrows: int, keep_raw: bool,
+                   allowed_ra: set[int] | None = None) -> int:
+    """Acquire one cover, distribute it, then reclaim its temporary raw table."""
+    raw_path = TMP_ROOT / "raw" / (cover.name + ".tbl")
+    if not raw_path.exists():
+        query_cover(cover, raw_path, outrows)
+    completed = process_cover(cover, raw_path, allowed_ra)
+    if not keep_raw:
+        raw_path.unlink()
+    print(f"completed={cover.name} final_tiles={completed}", flush=True)
+    return completed
+
+
+def reap_completed_covers(pending: dict[Future[int], tuple[Cover, Path]],
+                          keep_raw: bool) -> None:
+    """Collect at least one worker; only delete raw data after success."""
+    completed, _ = wait(pending, return_when=FIRST_COMPLETED)
+    for future in completed:
+        cover, raw_path = pending.pop(future)
+        final_tiles = future.result()
+        if not keep_raw:
+            raw_path.unlink(missing_ok=True)
+        print(f"processed={cover.name} final_tiles={final_tiles}", flush=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--plan", action="store_true", help="print cover count and exit")
+    parser.add_argument("--download", action="store_true", help="perform the IRSA acquisition")
+    parser.add_argument("--outrows", type=int, default=1_000_000,
+                        help="per-cone IRSA row cap; cap hits are fatal (default: %(default)s)")
+    parser.add_argument("--keep-raw", action="store_true", help="retain raw IPAC tables under /tmp")
+    parser.add_argument("--limit", type=int, default=None, help="only process this many unfinished covers")
+    parser.add_argument("--workers", type=int, default=2,
+                        help="parallel cleaning workers; downloads remain serial (default: %(default)s)")
+    args = parser.parse_args()
+    covers = [cover for dec_index in range(180) for cover in covers_for_band(dec_index)]
+    print(f"final_tiles=64800 covers={len(covers)} radius_deg={RADIUS_DEG:g} tmp_root={TMP_ROOT}")
+    print("estimated_final_cleaned_csv=15-17_GiB; raw/intermediate files are reclaimed per cover")
+    if args.plan:
+        return
+    if not args.download:
+        parser.error("choose --plan or --download")
+    if args.workers < 1:
+        parser.error("--workers must be at least 1")
+    if shutil.which("curl") is None:
+        parser.error("curl is required")
+    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
+    DONE_ROOT.mkdir(parents=True, exist_ok=True)
+    TMP_ROOT.mkdir(parents=True, exist_ok=True)
+    scheduled = 0
+    pending: dict[Future[int], tuple[Cover, Path]] = {}
+    # One serial producer downloads covers.  At most --workers downloaded
+    # covers await cleaning, bounding /tmp usage while overlapping the network
+    # wait with CPU-bound blackbody fitting and CSV distribution.
+    with ProcessPoolExecutor(max_workers=args.workers) as executor:
+        for cover in covers:
+            if not cover_has_work(cover):
+                continue
+            if args.limit is not None and scheduled >= args.limit:
+                break
+            while len(pending) >= args.workers:
+                reap_completed_covers(pending, args.keep_raw)
+            raw_path = TMP_ROOT / "raw" / (cover.name + ".tbl")
+            if not raw_path.exists():
+                query_cover(cover, raw_path, args.outrows)
+            print(f"downloaded={cover.name}; queued_cleaning={len(pending) + 1}/{args.workers}",
+                  flush=True)
+            pending[executor.submit(process_cover, cover, raw_path)] = (cover, raw_path)
+            scheduled += 1
+        while pending:
+            reap_completed_covers(pending, args.keep_raw)
+
+
+if __name__ == "__main__":
+    main()

scripts/process_2mass_psc.py

@@ -23,6 +23,7 @@ OUTPUT = ROOT / "assets/2mass/processed/2mass_psc_m31_0p5deg_stars.csv"
 WAVELENGTH_M = np.array([1.235, 1.662, 2.159]) * 1e-6
 ZERO_POINT_JY = np.array([1594.0, 1024.0, 666.7])
 GOOD_RD_FLAGS = frozenset("123")
+MISSING_NUMERIC_TOKENS = frozenset(("", "null"))
 MIN_TEMPERATURE_K = 300.0
 MAX_TEMPERATURE_K = 100000.0
 
@@ -62,7 +63,8 @@ def load_required_columns(path: Path) -> dict[str, np.ndarray]:
     result: dict[str, np.ndarray] = {}
     for name in ("ra", "dec", "j_m", "h_m", "k_m"):
         result[name] = np.array(
-            [float(value) if value else math.nan for value in values[name]], dtype=float
+            [float(value) if value.lower() not in MISSING_NUMERIC_TOKENS else math.nan
+             for value in values[name]], dtype=float
         )
     result["rd_flg"] = np.array(values["rd_flg"], dtype="U3")
     return result
@@ -123,11 +125,14 @@ def main() -> None:
     columns = load_required_columns(args.input)
     magnitudes = np.column_stack((columns["j_m"], columns["h_m"], columns["k_m"]))
     valid_photometry = np.isfinite(magnitudes).all(axis=1)
+    valid_position = (np.isfinite(columns["ra"]) & np.isfinite(columns["dec"]) &
+                      (columns["ra"] >= 0.0) & (columns["ra"] < 360.0) &
+                      (columns["dec"] >= -90.0) & (columns["dec"] <= 90.0))
     valid_rd_flag = np.array(
         [len(flag) == 3 and all(value in GOOD_RD_FLAGS for value in flag)
          for flag in columns["rd_flg"]]
     )
-    keep = valid_photometry & valid_rd_flag
+    keep = valid_position & valid_photometry & valid_rd_flag
     if not np.any(keep):
         raise ValueError("no sources retain valid J/H/Ks photometry")
 
@@ -145,6 +150,7 @@ def main() -> None:
                              f"{amplitude:.9e}"))
 
     print(f"input_rows={len(keep)}")
+    print(f"discarded_invalid_position={np.count_nonzero(~valid_position)}")
     print(f"discarded_invalid_photometry={np.count_nonzero(~valid_photometry)}")
     print(f"discarded_rd_flg_not_123={np.count_nonzero(~valid_rd_flag)}")
     print(f"retained_rows={np.count_nonzero(keep)}")