"""Sionna-compatible loader extended with geographic scene compilation."""
from __future__ import annotations
import json
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path
from typing import Any, Literal
from weakref import WeakKeyDictionary
from openworld_radio_twin.compiler import (
BuildingMode,
CellSize,
TerrainMode,
_add_derived_assets,
_geometry_request,
_run_sync,
_scene_lock,
cached_scene_directory,
compile_cached_scene,
resolve_scene_directory,
)
from openworld_radio_twin.geodesy import LocalFrame
from openworld_radio_twin.scene_format import SCENE_PROVENANCE_PATH, SCENE_XML_PATH
from openworld_radio_twin.simulation.scene_builder import (
MaterialProfileName,
load_scene_assets,
measurement_surface_path,
transmitter_scene_position,
)
from openworld_radio_twin.simulation.scene_builder import (
terrain_elevation as _terrain_elevation,
)
Device = Literal["auto", "cuda", "cpu"]
_SCENE_DIRECTORIES: WeakKeyDictionary[Any, Path] = WeakKeyDictionary()
__all__ = [
"Device",
"load_scene",
"load_scene_async",
"measurement_surface",
"position",
"scene_frame",
"scene_path",
"terrain_elevation",
]
def _version_tuple(value: str) -> tuple[int, int, int]:
numbers = []
for part in value.split(".")[:3]:
digits = "".join(character for character in part if character.isdigit())
numbers.append(int(digits or 0))
return tuple((*numbers, 0, 0)[:3])
def _sionna_runtime(device: Device):
if device not in {"auto", "cuda", "cpu"}:
raise ValueError("device must be 'auto', 'cuda', or 'cpu'")
try:
import mitsuba as mi
except ImportError as exc:
raise RuntimeError(
"Sionna RT is not installed; install 'openworld-radio-twin[rt]'"
) from exc
current = mi.variant()
requested = (
None if device == "auto" else f"{'cuda' if device == 'cuda' else 'llvm'}_ad_mono_polarized"
)
if current is None:
try:
if requested is None:
mi.set_variant("cuda_ad_mono_polarized", "llvm_ad_mono_polarized")
else:
mi.set_variant(requested)
except (ImportError, RuntimeError) as exc:
raise RuntimeError(
f"Could not initialize the requested Sionna {device} runtime"
) from exc
current = mi.variant()
elif requested is not None and current != requested:
raise RuntimeError(
f"Mitsuba variant is already '{current}' on this thread; device='{device}' requires "
f"'{requested}'. Select the device before importing sionna.rt, or use a separate "
"process."
)
if current not in {"cuda_ad_mono_polarized", "llvm_ad_mono_polarized"}:
raise RuntimeError(f"Mitsuba variant '{current}' is not compatible with Sionna RT")
import sionna.rt as sionna_rt
installed = getattr(sionna_rt, "__version__", None)
if not installed:
try:
installed = version("sionna-rt")
except PackageNotFoundError as exc:
raise RuntimeError(
"Sionna RT is not installed; install 'openworld-radio-twin[rt]'"
) from exc
if not isinstance(installed, str) or not installed.strip():
raise RuntimeError("The installed Sionna RT package has no readable version metadata")
if not (2, 0, 1) <= _version_tuple(installed) < (2, 1, 0):
raise RuntimeError(f"Sionna RT {installed} is unsupported; install version >=2.0.1,<2.1")
return sionna_rt
def _register(scene: Any, directory: Path | None) -> None:
if directory is None:
return
try:
_SCENE_DIRECTORIES[scene] = directory
except TypeError as exc:
raise RuntimeError("The installed Sionna Scene type cannot be weakly referenced") from exc
def _local_scene_source(filename: str | Path | None) -> tuple[str | None, Path | None]:
if filename is None:
return None, None
candidate = Path(filename).expanduser()
if candidate.is_dir():
xml_path = candidate / SCENE_XML_PATH
else:
xml_path = candidate
resolved = xml_path.resolve()
directory = resolved.parent if (resolved.parent / SCENE_PROVENANCE_PATH).is_file() else None
return str(resolved), directory
def _load_native_scene(
filename: str | Path | None,
merge_shapes: bool,
merge_shapes_exclude_regex: str | None,
remove_duplicate_vertices: bool,
device: Device,
sionna_options: dict[str, object],
):
sionna_rt = _sionna_runtime(device)
source, directory = _local_scene_source(filename)
scene = sionna_rt.load_scene(
source,
merge_shapes,
merge_shapes_exclude_regex,
remove_duplicate_vertices,
**sionna_options,
)
_register(scene, directory)
return scene
def _geographic_mode(
filename: str | Path | None,
latitude: float | None,
longitude: float | None,
radius_m: float | None,
) -> bool:
supplied = (latitude is not None, longitude is not None, radius_m is not None)
if any(supplied) and not all(supplied):
raise ValueError("latitude, longitude, and radius_m must be provided together")
if all(supplied) and filename is not None:
raise ValueError("filename and geographic coordinates are mutually exclusive")
return all(supplied)
[docs]
def load_scene(
filename: str | Path | None = None,
merge_shapes: bool = True,
merge_shapes_exclude_regex: str | None = None,
remove_duplicate_vertices: bool = False,
*,
latitude: float | None = None,
longitude: float | None = None,
radius_m: float | None = None,
cache_dir: str | Path | None = None,
buildings: BuildingMode = "auto",
building_source: str = "auto",
terrain: TerrainMode = "elevation",
material_profile: MaterialProfileName = "itu",
terrain_resolution_m: float = 5.0,
device: Device = "auto",
**sionna_options: object,
):
"""Load a native Sionna scene from a file, an empty scene, or WGS84 coordinates."""
if _geographic_mode(filename, latitude, longitude, radius_m):
assert latitude is not None and longitude is not None and radius_m is not None
filename = _run_sync(
compile_cached_scene(
latitude=latitude,
longitude=longitude,
radius_m=radius_m,
cache_dir=cache_dir,
buildings=buildings,
building_source=building_source,
terrain=terrain,
material_profile=material_profile,
terrain_resolution_m=terrain_resolution_m,
)
)
return _load_native_scene(
filename,
merge_shapes,
merge_shapes_exclude_regex,
remove_duplicate_vertices,
device,
sionna_options,
)
[docs]
async def load_scene_async(
*,
latitude: float,
longitude: float,
radius_m: float,
cache_dir: str | Path | None = None,
buildings: BuildingMode = "auto",
building_source: str = "auto",
terrain: TerrainMode = "elevation",
material_profile: MaterialProfileName = "itu",
terrain_resolution_m: float = 5.0,
device: Device = "auto",
merge_shapes: bool = True,
merge_shapes_exclude_regex: str | None = None,
remove_duplicate_vertices: bool = False,
**sionna_options: object,
):
"""Compile asynchronously, then load the native scene on the caller's event-loop thread."""
filename = await compile_cached_scene(
latitude=latitude,
longitude=longitude,
radius_m=radius_m,
cache_dir=cache_dir,
buildings=buildings,
building_source=building_source,
terrain=terrain,
material_profile=material_profile,
terrain_resolution_m=terrain_resolution_m,
)
return _load_native_scene(
filename,
merge_shapes,
merge_shapes_exclude_regex,
remove_duplicate_vertices,
device,
sionna_options,
)
[docs]
def scene_path(scene: Any) -> Path:
"""Return the canonical directory registered for a native Sionna scene."""
try:
return _SCENE_DIRECTORIES[scene]
except (KeyError, TypeError) as exc:
raise ValueError(
"Scene has no OWRT metadata; load it with openworld_radio_twin.rt.load_scene "
"or pass a scene directory to the geographic helper"
) from exc
def _directory(scene_or_path: Any) -> Path:
if isinstance(scene_or_path, (str, Path)):
return resolve_scene_directory(scene_or_path)
return scene_path(scene_or_path)
[docs]
def scene_frame(scene_or_path: Any) -> LocalFrame:
"""Restore the scene's public WGS84-to-ENU transform."""
directory = _directory(scene_or_path)
provenance = json.loads((directory / SCENE_PROVENANCE_PATH).read_text(encoding="utf-8"))
origin = provenance["coordinate_contract"]["origin"]
return LocalFrame.at(float(origin["longitude"]), float(origin["latitude"]))
[docs]
def terrain_elevation(scene_or_path: Any, east_m: float, north_m: float) -> float:
"""Sample local terrain elevation in metres at an ENU position."""
return _terrain_elevation(load_scene_assets(_directory(scene_or_path)), east_m, north_m)
[docs]
def position(
scene_or_path: Any,
*,
latitude: float,
longitude: float,
height_agl: float,
) -> list[float]:
"""Return the absolute ENU position expected by Sionna's radio devices."""
return list(
transmitter_scene_position(
load_scene_assets(_directory(scene_or_path)),
longitude,
latitude,
height_agl,
)
)
[docs]
def measurement_surface(
scene_or_path: Any,
*,
cell_size: CellSize = (10.0, 10.0),
height: float = 1.5,
):
"""Return a native Mitsuba mesh following terrain at a constant AGL height."""
directory = _directory(scene_or_path)
with _scene_lock(directory):
assets = load_scene_assets(directory)
_add_derived_assets(directory, assets, height, cell_size, None)
path = measurement_surface_path(assets, cell_size, height)
return _sionna_runtime("auto").load_mesh(str(path), flip_normals=False)
def cache_path(
*,
latitude: float,
longitude: float,
radius_m: float,
cache_dir: str | Path | None = None,
buildings: BuildingMode = "auto",
building_source: str = "auto",
terrain: TerrainMode = "elevation",
material_profile: MaterialProfileName = "itu",
terrain_resolution_m: float = 5.0,
) -> Path:
"""Resolve a geographic scene's deterministic cache directory without compiling it."""
request = _geometry_request(
latitude,
longitude,
radius_m,
buildings,
terrain,
material_profile,
terrain_resolution_m,
building_source,
)
return cached_scene_directory(request, cache_dir)
