Coverage for lunar_toolkit/config.py: 98%

1"""Flat YAML run configuration for ``lunar sfs``: load, validate, merge, resolve regions and dump."""

3from __future__ import annotations

5import dataclasses

6import difflib

7import hashlib

8import json

9import logging

10import os

11import types

12import typing

13from collections.abc import Mapping

14from dataclasses import dataclass, field

15from pathlib import Path

16from typing import Any, Union

18import yaml

20from .sfs import SfsConfig

22log = logging.getLogger(__name__)

24PGDA = "pgda"

25RUNS_DIR_ENV = "LUNAR_RUNS_DIR"

26"""``dem: pgda`` (or no ``dem`` with a ``region``): the region's PGDA LOLA 5 m ``surf`` DEM."""

28OPTION_KEYS = ("bundle_adjust_options", "refine_options", "sfs_options")

29FLOW_KEYS = ("projwin", "stereo_convergence", "crop_center")

32class ConfigError(ValueError):

33 """An invalid configuration: unknown key, wrong type or inconsistent settings."""

36@dataclass

37class RunConfig:

38 """Everything a ``lunar sfs`` run needs: inputs, region and the

39 [`SfsConfig`][lunar_toolkit.sfs.SfsConfig].

41 In YAML the run keys below and the [`SfsConfig`][lunar_toolkit.sfs.SfsConfig] fields share one flat

42 mapping. Relative paths are relative to the current directory.

44 Parameters

45 ----------

46 work_dir : Path

47 Run directory; a relative path is taken under *runs_dir*.

48 runs_dir : Path, optional

49 Where relative ``work_dir`` paths live (default: ``$LUNAR_RUNS_DIR``, else the

50 current directory).

51 Working / output directory.

52 dem : str, optional

53 Prior DEM path, or ``"pgda"`` for the region's PGDA LOLA 5 m DEM (the

54 default with a *region*).

55 region : str or int, optional

56 Artemis III region ([`lunar_toolkit.reference.get_region`][lunar_toolkit.reference.get_region]): sets

57 the ``projwin`` (its ROI, or a *crop_km* window), ``region_id`` and the

58 image search AOI.

59 crop_km : float, optional

60 Side of a square window at the region centre ([`crop_window`][lunar_toolkit.reference.crop_window]).

61 crop_center : (lat, lon), optional

62 Centre of that window (default: the ROI centre).

63 aoi : str, optional

64 AOI (GeoJSON / WKT file or WKT) for the image search and, without

65 ``projwin``, the crop.

66 sources : list of str

67 NAC product IDs or ``.IMG`` / ``.cub`` paths (default: search the AOI).

68 from_file : Path, optional

69 Text file of more sources, one per line.

70 max_slew, max_pixel_scale : float, optional

71 Image search cuts (deg, m/px). No slew cut by default: slewed images give the

72 stereo convergence for registration, and ``max_sfs_slew`` keeps them out of SfS.

73 max_images : int, optional

74 Cap on the searched images before download (stereo pairs + illumination coverage;

75 [`lunar_toolkit.sfs.cap_candidates`][lunar_toolkit.sfs.cap_candidates]).

76 pgda_dir : Path, optional

77 Where PGDA files are downloaded (default: the reference cache).

78 data_dir : Path, optional

79 Parquet index directory.

80 checksums : bool

81 MD5 the input files (under 2 GB) in ``provenance.json``.

82 sfs : SfsConfig

83 """

85 work_dir: Path = Path("sfs_work")

86 runs_dir: Path | None = None

87 dem: str | None = None

88 region: str | int | None = None

89 crop_km: float | None = None

90 crop_center: tuple[float, float] | None = None

91 aoi: str | None = None

92 sources: list[str] = field(default_factory=list)

93 from_file: Path | None = None

94 max_slew: float | None = None

95 max_pixel_scale: float | None = 1.5

96 max_images: int | None = None

97 pgda_dir: Path | None = None

98 data_dir: Path | None = None

99 checksums: bool = True

100 sfs: SfsConfig = field(default_factory=SfsConfig)

101

102

103def _hints(cls: type) -> dict[str, Any]:

104 names = [f.name for f in dataclasses.fields(cls)]

105 hints = typing.get_type_hints(cls)

106 return {n: hints[n] for n in names if n != "sfs"}

107

108

109RUN_KEYS: dict[str, Any] = _hints(RunConfig)

110SFS_KEYS: dict[str, Any] = _hints(SfsConfig)

111KEYS: dict[str, Any] = {**RUN_KEYS, **SFS_KEYS}

112

113

114def _describe(hint: Any) -> str:

115 if typing.get_origin(hint) in (Union, types.UnionType):

116 return " or ".join(_describe(a) for a in typing.get_args(hint))

117 origin, args = typing.get_origin(hint), typing.get_args(hint)

118 if origin is tuple:

119 return f"list of {len(args)} {_describe(args[0])}s"

120 if origin is list:

121 return f"list of {_describe(args[0])}s"

122 if origin is dict:

123 return "mapping"

124 names = {bool: "boolean", int: "integer", float: "number", str: "string", Path: "path"}

125 return "null" if hint is type(None) else names.get(hint, str(hint))

126

127

128def _coerce(key: str, value: Any, hint: Any) -> Any:

129 origin, args = typing.get_origin(hint), typing.get_args(hint)

130 if hint is Any: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 return value

132 if origin in (Union, types.UnionType):

133 if value is None and type(None) in args:

134 return None

135 for arm in args:

136 if arm is not type(None):

137 try:

138 return _coerce(key, value, arm)

139 except ConfigError:

140 pass

141 elif hint is bool and isinstance(value, bool):

142 return value

143 elif hint is int and isinstance(value, int) and not isinstance(value, bool):

144 return value

145 elif hint is float and isinstance(value, (int, float)) and not isinstance(value, bool):

146 return float(value)

147 elif hint is str and isinstance(value, str):

148 return value

149 elif hint is Path and isinstance(value, (str, Path)):

150 return Path(value)

151 elif origin is tuple and isinstance(value, (list, tuple)) and len(value) == len(args):

152 return tuple(_coerce(key, v, a) for v, a in zip(value, args))

153 elif origin is list and isinstance(value, (list, tuple)):

154 return [_coerce(key, v, args[0]) for v in value]

155 elif origin is dict and isinstance(value, Mapping) and all(isinstance(k, str) for k in value):

156 return {k.replace("-", "_"): v for k, v in value.items()}

157 raise ConfigError(f"{key}: expected {_describe(hint)}, got {value!r}")

158

159

160def from_dict(data: Mapping[str, Any], source: str = "config") -> RunConfig:

161 """A validated [`RunConfig`][lunar_toolkit.config.RunConfig] from a flat mapping (keys as in `KEYS`).

162

163 Raises

164 ------

165 ConfigError

166 For unknown keys (with suggestions), wrong types and invalid choices.

167 """

168 if not isinstance(data, Mapping):

169 raise ConfigError(f"{source}: expected a mapping of settings, got {type(data).__name__}")

170 norm = {str(k).replace("-", "_"): v for k, v in data.items()}

171 for key in norm:

172 if key not in KEYS:

173 close = difflib.get_close_matches(key, list(KEYS), n=1)

174 hint = f" (did you mean {close[0]!r}?)" if close else ""

175 raise ConfigError(f"{source}: unknown key {key!r}{hint}")

176 try:

177 values = {k: _coerce(k, v, KEYS[k]) for k, v in norm.items()}

178 except ConfigError as exc:

179 raise ConfigError(f"{source}: {exc}") from None

180 try:

181 sfs = SfsConfig(**{k: v for k, v in values.items() if k in SFS_KEYS})

182 except ValueError as exc:

183 raise ConfigError(f"{source}: {exc}") from None

184 return RunConfig(**{k: v for k, v in values.items() if k in RUN_KEYS}, sfs=sfs)

185

186

187def _plain(value: Any) -> Any:

188 if isinstance(value, Path):

189 return str(value)

190 if isinstance(value, (list, tuple)):

191 return [_plain(v) for v in value]

192 if isinstance(value, dict):

193 return {k: _plain(v) for k, v in value.items()}

194 return value

195

196

197def to_dict(cfg: RunConfig) -> dict[str, Any]:

198 """Flat, YAML / JSON-ready mapping of every setting (run keys first)."""

199 out = {k: _plain(getattr(cfg, k)) for k in RUN_KEYS}

200 out.update({k: _plain(getattr(cfg.sfs, k)) for k in SFS_KEYS})

201 return out

202

203

204def merge(cfg: RunConfig, overrides: Mapping[str, Any], source: str = "overrides") -> RunConfig:

205 """*cfg* with *overrides* applied; option dicts are merged key by key, other values replaced."""

206 data = to_dict(cfg)

207 for key, value in overrides.items():

208 key = key.replace("-", "_")

209 if key in OPTION_KEYS and isinstance(value, Mapping) and isinstance(data.get(key), dict):

210 data[key] = {**data[key], **value}

211 else:

212 data[key] = value

213 return from_dict(data, source)

214

215

216def load_config(path: Path) -> RunConfig:

217 """Read and validate a YAML config (an empty file gives the defaults)."""

218 path = Path(path)

219 try:

220 data = yaml.safe_load(path.read_text())

221 except OSError as exc:

222 raise ConfigError(f"Cannot read config {path}: {exc.strerror or exc}") from None

223 except yaml.YAMLError as exc:

224 raise ConfigError(f"{path}: invalid YAML: {exc}") from None

225 return from_dict(data if data is not None else {}, str(path))

226

227

228def dump_config(cfg: RunConfig, path: Path | None = None) -> str:

229 """YAML text of every setting (written to *path* when given)."""

230 from .constants import APP_VERSION

231

232 lines = [f"# lunar sfs configuration (lunar-toolkit {APP_VERSION})\n"]

233 for key, value in to_dict(cfg).items():

234 if key == next(iter(SFS_KEYS)):

235 lines.append("\n# SfS pipeline settings (lunar_toolkit.sfs.SfsConfig)\n")

236 if key in FLOW_KEYS or (isinstance(value, (list, dict)) and not value):

237 flow = yaml.safe_dump(value, default_flow_style=True, width=10**6).strip()

238 lines.append(f"{key}: {flow.removesuffix('...').strip()}\n")

239 else:

240 lines.append(yaml.safe_dump({key: value}, default_flow_style=False, sort_keys=False, width=110))

241 text = "".join(lines)

242 if path is not None:

243 Path(path).parent.mkdir(parents=True, exist_ok=True)

244 Path(path).write_text(text)

245 return text

246

247

248def config_hash(cfg: RunConfig) -> str:

249 """SHA-256 of the settings (canonical JSON of [`to_dict`][lunar_toolkit.config.to_dict])."""

250 return hashlib.sha256(json.dumps(to_dict(cfg), sort_keys=True).encode()).hexdigest()

251

252

253def pgda_file(region: Any, kind: str, out_dir: Path | None = None, download: bool = True) -> Path:

254 """Local path of a region's PGDA LOLA 5 m raster (``surf``, ``ldec``, ``toterr``, ...).

255

256 Downloaded into ``out_dir/<site>/`` when missing and *download*; otherwise the

257 expected path (``<site>_final_adj_5mpp_<kind>.tif``, ``<site>_final_adj.xyzi``)

258 is returned without touching the network or the disk.

259 """

260 from . import reference as ref

261

262 site = ref._site_name(region)

263 out_dir = Path(out_dir) if out_dir else ref.REFERENCE_CACHE_DIR / "pgda"

264 name = f"{site}_final_adj.xyzi" if kind == "xyzi" else f"{site}_final_adj_5mpp_{kind}.tif"

265 expected = out_dir / site / name

266 if expected.exists() or not download:

267 return expected

268 files = ref.download_pgda(site, [kind], out_dir)[kind]

269 if not files: 269 ↛ 270line 269 didn't jump to line 270 because the condition on line 269 was never true

270 raise ConfigError(f"PGDA site {site} has no {kind} file")

271 return files[0]

272

273

274def region_id(region: Any) -> str:

275 """Bertone et al. product prefix of a region, e.g. ``A311`` for Malapert Massif."""

276 from . import reference as ref

277

278 return f"A3{ref.get_region(region).id:02d}"

279

280

281def resolve(cfg: RunConfig, download: bool = True) -> RunConfig:

282 """Fill what *region* implies: ``projwin``, ``region_id`` and the PGDA prior DEM.

283

284 An explicit ``projwin`` wins, then ``aoi`` (the crop is computed from it at run

285 time), then a ``crop_km`` window, then the region ROI. ``dem: pgda`` (or none)

286 is the region's PGDA ``surf`` DEM, downloaded when missing and *download*.

287

288 Raises

289 ------

290 ConfigError

291 For a ``crop_km`` / ``dem: pgda`` without a region, or an unknown region.

292 """

293 from . import reference as ref

294

295 runs_dir = cfg.runs_dir or (Path(os.environ[RUNS_DIR_ENV]) if os.environ.get(RUNS_DIR_ENV) else None)

296 if runs_dir is not None and not Path(cfg.work_dir).is_absolute():

297 cfg = merge(cfg, {"work_dir": str(Path(runs_dir).expanduser() / cfg.work_dir), "runs_dir": None})

298 if cfg.region is None:

299 if cfg.crop_km is not None or cfg.crop_center is not None:

300 raise ConfigError("crop_km / crop_center need a region")

301 if cfg.dem == PGDA:

302 raise ConfigError("dem: pgda needs a region")

303 return cfg

304 try:

305 reg = ref.get_region(cfg.region)

306 except KeyError as exc:

307 raise ConfigError(f"region: {exc.args[0]}") from None

308 changes: dict[str, Any] = {}

309 if cfg.sfs.projwin is None and cfg.aoi is None:

310 if cfg.crop_km is not None:

311 box = ref.crop_window(reg, cfg.crop_km, cfg.crop_center)

312 else:

313 box = reg.roi

314 changes["projwin"] = box.projwin

315 if cfg.sfs.region_id is None:

316 changes["region_id"] = region_id(reg)

317 if cfg.dem in (None, PGDA):

318 try:

319 changes["dem"] = str(pgda_file(reg, "surf", cfg.pgda_dir, download=download))

320 except KeyError as exc:

321 raise ConfigError(f"dem: {exc.args[0]}") from None

322 return merge(cfg, changes, "region") if changes else cfg

323

324

325def search_aoi(cfg: RunConfig) -> str | None:

326 """Image-search AOI: ``aoi``, else the region crop (``projwin`` as lon/lat WKT), else None."""

327 from . import reference as ref

328

329 if cfg.aoi is not None:

330 return cfg.aoi

331 if cfg.region is None or cfg.sfs.projwin is None:

332 return None

333 ulx, uly, lrx, lry = cfg.sfs.projwin

334 return ref.PsBox(ulx, lry, lrx, uly).wkt()

335

336

337def load_sources(cfg: RunConfig) -> list[str]:

338 """``sources`` plus the non-comment lines of ``from_file``."""

339 sources = list(cfg.sources)

340 if cfg.from_file is not None:

341 try:

342 text = Path(cfg.from_file).read_text()

343 except OSError as exc:

344 raise ConfigError(f"from_file: cannot read {cfg.from_file}: {exc.strerror or exc}") from None

345 sources += [s.strip() for s in text.splitlines() if s.strip() and not s.strip().startswith("#")]

346 return sources