#!/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()