280 lines
12 KiB
Python
280 lines
12 KiB
Python
#!/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()
|