Coverage for lunar_toolkit/kaguya/tc.py: 76%
192 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 02:58 +0000
1"""Kaguya Terrain Camera Level 2B0: ODE file selection, download and ISIS import."""
3from __future__ import annotations
5import gzip
6import logging
7import re
8import shutil
9from collections.abc import Iterable, Sequence
10from concurrent.futures import ThreadPoolExecutor, as_completed
11from dataclasses import dataclass
12from pathlib import Path
13from typing import TYPE_CHECKING, Any
15from ..constants import DEFAULT_PROJECTION, GEOTIFF_OPTIONS
16from ..downloader import download_file, url_filename
17from ..isis import isd_generate, kaguyatc2isis, map_project, spiceinit
18from ..tools import is_dry_run, reuse_existing, run
20if TYPE_CHECKING: # pragma: no cover
21 import pandas as pd
23 from ..ode import OdeClient
25log = logging.getLogger(__name__)
27# Level 2B0 TC product IDs, e.g. TC1W2B0_01_02936N034E0938 (W = stereo, S = mono).
28TC_ID_RE = re.compile(r"^TC[12][WS]2B0_\d{2}_\d{5}[NS]\d{3}E\d{4}$", re.IGNORECASE)
30# kaguyatc2isis: keep every DN (ASP turns the special-pixel ranges off).
31KAGUYATC2ISIS_PARAMS: dict[str, Any] = {
32 "setnullrange": False,
33 "sethrsrange": False,
34 "sethisrange": False,
35 "setlrsrange": False,
36 "setlisrange": False,
37}
39# spiceinit: reconstructed CK, smithed / reconstructed SPK (ASP example).
40# ``attach`` is dropped with ``web=true``, which always attaches (ISIS rejects both).
41SPICEINIT_PARAMS: dict[str, Any] = {
42 "attach": True,
43 "cksmithed": False,
44 "ckrecon": True,
45 "ckpredicted": False,
46 "cknadir": False,
47 "spkrecon": True,
48 "spkpredicted": False,
49 "shape": "system",
50 "startpad": 0.0,
51 "endpad": 0.0,
52}
54_IMAGE_POINTER_RE = re.compile(r'^\s*\^IMAGE\s*=\s*\(?\s*"([^"]+)"', re.MULTILINE)
57@dataclass(frozen=True)
58class TcProduct:
59 """Data and label URLs of one Level 2B0 TC product."""
61 product_id: str
62 image_url: str
63 label_url: str
64 image_kbytes: float | None = None
66 @property
67 def is_stereo(self) -> bool:
68 return self.product_id[3] == "W"
70 @property
71 def sensor(self) -> str:
72 """``"TC1"`` or ``"TC2"``."""
73 return self.product_id[:3]
76@dataclass
77class TcOutputs:
78 """Files written for one product (None for skipped steps)."""
80 product_id: str
81 cube: Path
82 camera: Path | None = None
83 map_cube: Path | None = None
84 geotiff: Path | None = None
87def is_tc_product_id(product_id: str) -> bool:
88 """True for a Level 2B0 TC product ID that ``kaguyatc2isis`` can import."""
89 return bool(TC_ID_RE.match(product_id.strip()))
92def select_tc_files(files: Iterable[dict], product_id: str | None = None) -> TcProduct | None:
93 """Pick the data and label files of a TC product from its ODE ``files`` list.
95 Parameters
96 ----------
97 files : iterable of dict
98 ODE product-file dicts (``FileName``, ``Type``, ``URL``, ``KBytes``).
99 product_id : str, optional
100 Product ID; derived from the data file name when omitted.
102 Returns
103 -------
104 TcProduct or None
105 None when the list lacks a ``.img`` / ``.img.gz`` data file or a
106 ``.lbl`` label. An uncompressed ``.img`` is preferred over ``.img.gz``.
107 """
108 image: dict | None = None
109 label: dict | None = None
110 for f in files or []:
111 if str(f.get("Type", "")).lower() != "product" or not f.get("URL"):
112 continue
113 name = str(f.get("FileName") or url_filename(f["URL"])).upper()
114 if name.endswith(".LBL") and label is None:
115 label = f
116 elif name.endswith(".IMG") or (name.endswith(".IMG.GZ") and image is None): 116 ↛ 110line 116 didn't jump to line 110 because the condition on line 116 was always true
117 image = f
118 if image is None or label is None:
119 return None
120 pid = product_id or url_filename(image["URL"]).split(".")[0]
121 kb = image.get("KBytes")
122 try:
123 kbytes = float(kb) if kb is not None else None
124 except (TypeError, ValueError):
125 kbytes = None
126 return TcProduct(pid.upper(), image["URL"], label["URL"], kbytes)
129def tc_products(df: pd.DataFrame) -> list[TcProduct]:
130 """TC products with downloadable data + label files from an ODE search.
132 Rows that are not Level 2B0 products (DTMs, ortho maps) or lack either
133 file are skipped with a warning.
134 """
135 out: list[TcProduct] = []
136 pids = df["product_id"] if "product_id" in df else [None] * len(df)
137 for pid, files in zip(pids, df.get("files", [])):
138 if pid and not is_tc_product_id(pid):
139 log.warning("Skipping %s: not a Level 2B0 TC product", pid)
140 continue
141 product = select_tc_files(files, pid)
142 if product is None: 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true
143 log.warning("Skipping %s: no .img/.lbl product files", pid)
144 continue
145 out.append(product)
146 return out
149def lookup_products(
150 product_ids: Iterable[str],
151 client: OdeClient | None = None,
152) -> pd.DataFrame:
153 """Query ODE by product ID and return the observation DataFrame (see
154 [`OdeClient.lookup`][lunar_toolkit.ode.OdeClient.lookup])."""
155 from ..ode import OdeClient
157 return (client or OdeClient()).lookup(product_ids, instrument_label="Kaguya TC")
160def image_pointer(label: Path) -> str | None:
161 """File name in the label's ``^IMAGE`` pointer (None if attached / absent)."""
162 m = _IMAGE_POINTER_RE.search(Path(label).read_text(encoding="latin-1", errors="replace"))
163 return m.group(1) if m else None
166def decompress(gz: Path, out: Path | None = None, remove: bool = False) -> Path:
167 """Gunzip *gz* to *out* (default: name without ``.gz``) via a ``.part`` file."""
168 gz = Path(gz)
169 out = Path(out) if out else gz.with_suffix("")
170 tmp = out.with_name(out.name + ".part")
171 log.info("Decompressing %s", gz.name)
172 with gzip.open(gz, "rb") as src, open(tmp, "wb") as dst:
173 shutil.copyfileobj(src, dst, 8 << 20)
174 tmp.replace(out)
175 if remove: 175 ↛ 177line 175 didn't jump to line 177 because the condition on line 175 was always true
176 gz.unlink()
177 return out
180def download_product(
181 product: TcProduct,
182 output_dir: Path,
183 force: bool = False,
184 keep_gz: bool = False,
185 show_progress: bool = True,
186) -> Path:
187 """Download (and gunzip) one product's label and image into *output_dir*.
189 The image is written under the name in the label's ``^IMAGE`` pointer so
190 ``kaguyatc2isis`` finds it next to the label. Existing files are reused
191 unless *force*. Returns the label path.
192 """
193 output_dir = Path(output_dir)
194 output_dir.mkdir(parents=True, exist_ok=True)
195 label = output_dir / url_filename(product.label_url)
196 if not reuse_existing(label, force):
197 download_file(product.label_url, label, show_progress=show_progress)
198 img = output_dir / (image_pointer(label) or f"{product.product_id}.img")
199 if reuse_existing(img, force):
200 return label
201 remote = output_dir / url_filename(product.image_url)
202 if not reuse_existing(remote, force): 202 ↛ 204line 202 didn't jump to line 204 because the condition on line 202 was always true
203 download_file(product.image_url, remote, show_progress=show_progress)
204 if remote.name.lower().endswith(".gz"): 204 ↛ 206line 204 didn't jump to line 206 because the condition on line 204 was always true
205 decompress(remote, img, remove=not keep_gz)
206 elif remote != img:
207 remote.replace(img)
208 return label
211class TcProcessor:
212 """Download and import Kaguya TC Level 2B0 products into ISIS.
214 Each product ends up in *output_dir* as ``<id>.cub`` (spiceinit'ed),
215 ``<id>.json`` (CSM camera, when *csm*) and, when *mapproject*,
216 ``<id>.tif`` (`GEOTIFF_OPTIONS`), the
217 same flat layout and resume / *keep* / *force* behaviour as
218 [`NacProcessor`][lunar_toolkit.lroc.nac.NacProcessor]. Steps whose output
219 exists are skipped unless *force*.
221 Parameters
222 ----------
223 output_dir : Path
224 Where labels, images, cubes and cameras are written.
225 web : bool
226 Use the USGS SPICE web service in ``spiceinit`` (no local Kaguya
227 kernels needed). Set False when ``$ISISDATA/kaguya`` is installed.
228 csm : bool
229 Write a CSM camera with ``isd_generate``.
230 only_naif_spice : bool, optional
231 Pass ``--only_naif_spice`` to ``isd_generate`` (ASP's recipe; needs
232 local kernels). Default: ``not web``.
233 mapproject : bool
234 Also ``cam2map`` the cube and write a GeoTIFF.
235 projection : str
236 ISIS projection for ``mosrange`` (default: equirectangular).
237 resolution : float, optional
238 Map resolution in m/pixel (default: ``mosrange``'s estimate).
239 keep : bool
240 Keep the PDS ``.img`` and the map-projected cube after use.
241 force : bool
242 Redo every step even when its output exists.
243 kaguyatc2isis_params, spiceinit_params : dict, optional
244 Overrides merged into `KAGUYATC2ISIS_PARAMS` /
245 `SPICEINIT_PARAMS`.
246 """
248 def __init__(
249 self,
250 output_dir: Path = Path("kaguya"),
251 *,
252 web: bool = True,
253 csm: bool = True,
254 only_naif_spice: bool | None = None,
255 mapproject: bool = False,
256 projection: str = DEFAULT_PROJECTION,
257 resolution: float | None = None,
258 keep: bool = False,
259 force: bool = False,
260 kaguyatc2isis_params: dict[str, Any] | None = None,
261 spiceinit_params: dict[str, Any] | None = None,
262 ) -> None:
263 self.output_dir = Path(output_dir)
264 self.web = web
265 self.csm = csm
266 self.only_naif_spice = (not web) if only_naif_spice is None else only_naif_spice
267 self.mapproject = mapproject
268 self.projection = projection
269 self.resolution = resolution
270 self.keep = keep
271 self.force = force
272 self.kaguyatc2isis_params = {**KAGUYATC2ISIS_PARAMS, **(kaguyatc2isis_params or {})}
273 self.spiceinit_params = {**SPICEINIT_PARAMS, **(spiceinit_params or {})}
275 def process(self, source: TcProduct | str | Path) -> TcOutputs:
276 """Download (if needed) and import one product.
278 *source* is a [`TcProduct`][lunar_toolkit.kaguya.tc.TcProduct], a local ``.lbl`` path or a product
279 ID (looked up on ODE).
280 """
281 if isinstance(source, TcProduct):
282 label = download_product(source, self.output_dir, force=self.force)
283 elif Path(source).suffix.lower() == ".lbl":
284 label = Path(source)
285 else:
286 products = tc_products(lookup_products([str(source)]))
287 if not products:
288 raise ValueError(f"No Kaguya TC Level 2B0 product {source!r} on ODE")
289 label = download_product(products[0], self.output_dir, force=self.force)
290 return self.import_label(label)
292 def process_batch(
293 self,
294 sources: Sequence[TcProduct | str | Path],
295 max_workers: int = 2,
296 ) -> list[TcOutputs | None]:
297 """Process several products in parallel (None for failures)."""
298 results: list[TcOutputs | None] = [None] * len(sources)
299 with ThreadPoolExecutor(max_workers=max_workers) as pool:
300 futures = {pool.submit(self.process, s): i for i, s in enumerate(sources)}
301 for future in as_completed(futures):
302 i = futures[future]
303 try:
304 results[i] = future.result()
305 except Exception as exc:
306 log.error("Failed %s: %s", _source_name(sources[i]), exc)
307 return results
309 def import_label(self, label: Path) -> TcOutputs:
310 """Run the ISIS chain on a downloaded PDS3 label (image beside it)."""
311 label = Path(label)
312 pid = label.name.split(".")[0].upper()
313 if not is_dry_run():
314 self.output_dir.mkdir(parents=True, exist_ok=True)
315 cube = self.output_dir / f"{pid}.cub"
316 out = TcOutputs(pid, cube)
318 if not reuse_existing(cube, self.force):
319 try:
320 kaguyatc2isis(label, cube, **self.kaguyatc2isis_params)
321 spiceinit(cube, web=self.web, **self._spiceinit_params())
322 except Exception:
323 cube.unlink(missing_ok=True) # never reuse a cube without SPICE
324 raise
325 if not self.keep and not is_dry_run(): 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 self._remove_image(label)
328 if self.csm: 328 ↛ 333line 328 didn't jump to line 333 because the condition on line 328 was always true
329 out.camera = self.output_dir / f"{pid}.json"
330 if not reuse_existing(out.camera, self.force):
331 isd_generate(cube, out.camera, only_naif_spice=self.only_naif_spice)
333 if self.mapproject:
334 out.geotiff = self.output_dir / f"{pid}.tif"
335 map_cube = self.output_dir / f"{pid}.map.cub"
336 if not reuse_existing(out.geotiff, self.force): 336 ↛ 338line 336 didn't jump to line 338 because the condition on line 336 was always true
337 self._mapproject(cube, map_cube, out.geotiff)
338 if self.keep: 338 ↛ 339line 338 didn't jump to line 339 because the condition on line 338 was never true
339 out.map_cube = map_cube
340 return out
342 def _mapproject(self, cube: Path, map_cube: Path, tif: Path) -> None:
343 map_project(cube, map_cube, self.projection, self.resolution, keep=self.keep)
344 run("gdal_translate", [*GEOTIFF_OPTIONS, map_cube, tif])
345 if not self.keep and not is_dry_run(): 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 map_cube.unlink(missing_ok=True)
348 def _spiceinit_params(self) -> dict[str, Any]:
349 params = dict(self.spiceinit_params)
350 if self.web:
351 params.pop("attach", None)
352 return params
354 @staticmethod
355 def _remove_image(label: Path) -> None:
356 try:
357 name = image_pointer(label)
358 except OSError:
359 return
360 if name:
361 (label.parent / name).unlink(missing_ok=True)
364def import_tc(label: Path, output_dir: Path | None = None, **kwargs: Any) -> TcOutputs:
365 """Import one downloaded TC product; see [`TcProcessor`][lunar_toolkit.kaguya.tc.TcProcessor] for
366 *kwargs*."""
367 label = Path(label)
368 return TcProcessor(output_dir or label.parent, **kwargs).import_label(label)
371def _source_name(source: Any) -> str:
372 return source.product_id if isinstance(source, TcProduct) else str(source)