zero2geospatial3d-sdk 0.1.0__py3-none-any.whl

zero2geospatial3d/scene.py

@@ -0,0 +1,555 @@
+"""Headless 3D scene: terrain + extruded buildings + trees + roads.
+
+A :class:`Scene` is the SDK's equivalent of the plugin's exported scene. It
+rasterises footprints onto the terrain grid (giving the height *surface* used
+by :mod:`zero2geospatial3d.shading` and the per-cell building ids used for
+roof statistics), and writes the same package the browser viewer reads
+(``scene.json``, ``terrain.bin``, ``L*.geojson``), so a scene built in Python
+opens in :func:`zero2geospatial3d.server.serve_scene`.
+
+Coordinates are local metres: x east, y north, origin at the scene centre.
+"""
+
+from __future__ import annotations
+
+import datetime as _dt
+import json
+import math
+import os
+from dataclasses import dataclass, field
+from typing import Any, Optional, Sequence
+
+import numpy as np
+
+from . import geo
+from .grid import Grid
+
+__all__ = [
+    "SCENE_FORMAT",
+    "SCENE_VERSION",
+    "Building",
+    "Tree",
+    "Road",
+    "Scene",
+    "rasterize_polygons",
+    "load_scene",
+]
+
+SCENE_FORMAT = "zero2geospatial3d-scene"
+SCENE_VERSION = 1
+
+Ring = Sequence[Sequence[float]]
+
+
+@dataclass
+class Building:
+    """An extruded footprint. ``polygons`` is a list of polygons, each a list of
+    rings (outer first, holes after) of local ``(x, y)``."""
+
+    polygons: list[list[Ring]]
+    height: float = 9.0
+    base: float = 0.0
+    id: int = 0
+    color: str = "#d9d4cc"
+    properties: dict[str, Any] = field(default_factory=dict)
+    group: Optional[str] = None
+
+
+@dataclass
+class Tree:
+    """A tree as a cylinder crown (``crown_radius``) from ``crown_base`` to ``height``."""
+
+    x: float
+    y: float
+    height: float = 8.0
+    crown_radius: float = 3.0
+    crown_base: float = 2.5
+    properties: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class Road:
+    """A polyline traffic source with a relative ``traffic_weight``."""
+
+    coords: list[Sequence[float]]
+    traffic_weight: float = 1.0
+    properties: dict[str, Any] = field(default_factory=dict)
+
+
+def _ring_mask(xs: np.ndarray, ys: np.ndarray, ring: Ring) -> np.ndarray:
+    """Even-odd point-in-ring for arrays of points."""
+    inside = np.zeros(xs.shape, dtype=bool)
+    pts = np.asarray(ring, dtype=float)[:, :2]
+    n = len(pts)
+    j = n - 1
+    for i in range(n):
+        xi, yi = pts[i]
+        xj, yj = pts[j]
+        if yi != yj:
+            crosses = (yi > ys) != (yj > ys)
+            xint = (xj - xi) * (ys - yi) / (yj - yi) + xi
+            inside ^= crosses & (xs < xint)
+        j = i
+    return inside
+
+
+def rasterize_polygons(
+    grid: Grid, polygons: Sequence[list[Ring]], values: Optional[np.ndarray] = None
+) -> np.ndarray:
+    """Boolean mask of cells whose centre falls inside any of ``polygons``.
+
+    Polygons smaller than a cell still mark the cell under their first vertex.
+    """
+    mask = np.zeros(grid.shape, dtype=bool)
+    for poly in polygons:
+        if not poly or len(poly[0]) < 3:
+            continue
+        outer = np.asarray(poly[0], dtype=float)
+        c0 = max(0, int(math.floor((outer[:, 0].min() - grid.min_x) / grid.cell)))
+        c1 = min(grid.cols - 1, int(math.ceil((outer[:, 0].max() - grid.min_x) / grid.cell)))
+        r0 = max(0, int(math.floor((grid.max_y - outer[:, 1].max()) / grid.cell)))
+        r1 = min(grid.rows - 1, int(math.ceil((grid.max_y - outer[:, 1].min()) / grid.cell)))
+        if c0 > c1 or r0 > r1:
+            continue
+        xs = grid.min_x + np.arange(c0, c1 + 1) * grid.cell
+        ys = grid.max_y - np.arange(r0, r1 + 1) * grid.cell
+        xx, yy = np.meshgrid(xs, ys)
+        inside = _ring_mask(xx, yy, poly[0])
+        for hole in poly[1:]:
+            inside &= ~_ring_mask(xx, yy, hole)
+        if not inside.any():
+            c, r = grid.index_of(float(outer[0, 0]), float(outer[0, 1]))
+            if c0 <= c <= c1 and r0 <= r <= r1:
+                inside[r - r0, c - c0] = True
+        mask[r0 : r1 + 1, c0 : c1 + 1] |= inside
+    return mask
+
+
+class Scene:
+    """Terrain grid plus objects, with the derived analysis rasters.
+
+    Args:
+        terrain: DEM on the analysis grid (local metres). Use
+            :meth:`Scene.flat` for a level site.
+        buildings, trees, roads: scene objects in local metres.
+        lat, lon: WGS84 position of the local origin (for the sun).
+        utc_offset: clock time zone; defaults to the nautical zone of ``lon``.
+    """
+
+    def __init__(
+        self,
+        terrain: Grid,
+        buildings: Optional[list[Building]] = None,
+        trees: Optional[list[Tree]] = None,
+        roads: Optional[list[Road]] = None,
+        *,
+        lat: float = 0.0,
+        lon: float = 0.0,
+        utc_offset: Optional[float] = None,
+        title: str = "Scene",
+        crs: Optional[dict[str, str]] = None,
+        origin: tuple[float, float] = (0.0, 0.0),
+    ) -> None:
+        self.terrain = terrain
+        self.buildings = list(buildings or [])
+        self.trees = list(trees or [])
+        self.roads = list(roads or [])
+        self.lat = lat
+        self.lon = lon
+        self.utc_offset = geo.utc_offset_guess(lon) if utc_offset is None else float(utc_offset)
+        self.title = title
+        self.crs = crs or {"authid": "", "wkt": ""}
+        self.origin = origin
+        self._number_buildings()
+
+    def _number_buildings(self) -> None:
+        """Give buildings without an id the next free positive id."""
+        used = {b.id for b in self.buildings if b.id}
+        nxt = max(used, default=0)
+        for b in self.buildings:
+            if not b.id:
+                nxt += 1
+                b.id = nxt
+
+    # ------------------------------------------------------------ builders
+    @classmethod
+    def flat(
+        cls, width: float, height: float, cell: float = 2.0, elevation: float = 0.0, **kwargs: Any
+    ) -> Scene:
+        """A level site of ``width`` x ``height`` metres centred on the origin."""
+        cols = max(2, int(math.ceil(width / cell)))
+        rows = max(2, int(math.ceil(height / cell)))
+        grid = Grid(
+            np.full((rows, cols), float(elevation)),
+            cell,
+            -cols * cell / 2 + cell / 2,
+            rows * cell / 2 - cell / 2,
+        )
+        return cls(grid, **kwargs)
+
+    @classmethod
+    def from_osm(
+        cls,
+        features: dict[str, list[dict[str, Any]]],
+        *,
+        lat: float,
+        lon: float,
+        half_size: float = 250.0,
+        cell: float = 2.0,
+        terrain: Optional[Grid] = None,
+        **kwargs: Any,
+    ) -> Scene:
+        """Scene from :func:`zero2geospatial3d.osm.payload_to_features` output.
+
+        Coordinates are projected around ``(lon, lat)`` with
+        :func:`zero2geospatial3d.geo.lonlat_to_local`.
+        """
+
+        def local(pt: Sequence[float]) -> list[float]:
+            x, y = geo.lonlat_to_local(pt[0], pt[1], lon, lat)
+            return [round(x, 2), round(y, 2)]
+
+        scene = terrain
+        if scene is None:
+            base = cls.flat(2 * half_size, 2 * half_size, cell)
+            scene = base.terrain
+        buildings = []
+        for k, f in enumerate(features.get("buildings", []), start=1):
+            g = f["geometry"]
+            polys = [g["coordinates"]] if g["type"] == "Polygon" else g["coordinates"]
+            p = f.get("properties", {})
+            buildings.append(
+                Building(
+                    [[[local(pt) for pt in ring] for ring in poly] for poly in polys],
+                    height=float(p.get("height_m", 9.0)),
+                    base=float(p.get("min_height_m", 0.0)),
+                    id=k,
+                    properties=dict(p),
+                    group=p.get("bldg_group"),
+                )
+            )
+        trees = []
+        for f in features.get("trees", []):
+            x, y = local(f["geometry"]["coordinates"])
+            h = float(f.get("properties", {}).get("height_m", 8.0))
+            trees.append(
+                Tree(
+                    x,
+                    y,
+                    h,
+                    crown_radius=max(1.5, 0.35 * h),
+                    crown_base=0.3 * h,
+                    properties=dict(f.get("properties", {})),
+                )
+            )
+        roads = []
+        for f in features.get("roads", []):
+            p = f.get("properties", {})
+            roads.append(
+                Road(
+                    [local(pt) for pt in f["geometry"]["coordinates"]],
+                    float(p.get("traffic_weight", 1.0)),
+                    dict(p),
+                )
+            )
+        kwargs.setdefault("title", "OSM scene")
+        return cls(scene, buildings, trees, roads, lat=lat, lon=lon, **kwargs)
+
+    # ------------------------------------------------------------ rasters
+    def building_rasters(self) -> tuple[np.ndarray, np.ndarray]:
+        """``(height_above_ground, ids)`` per cell; id 0 is open ground.
+
+        Taller buildings win where footprints overlap. Raised parts
+        (``base > 0``) are treated as solid down to the ground (2.5-D).
+        """
+        self._number_buildings()
+        heights = np.zeros(self.terrain.shape)
+        ids = np.zeros(self.terrain.shape, dtype=np.int32)
+        for b in sorted(self.buildings, key=lambda b: b.height):
+            m = rasterize_polygons(self.terrain, b.polygons)
+            heights[m] = np.maximum(heights[m], b.height)
+            ids[m] = b.id
+        return heights, ids
+
+    def tree_raster(self) -> tuple[np.ndarray, np.ndarray]:
+        """``(crown_top_above_ground, canopy_mask)`` from the tree list."""
+        top = np.zeros(self.terrain.shape)
+        xx, yy = self.terrain.coords()
+        for t in self.trees:
+            m = (xx - t.x) ** 2 + (yy - t.y) ** 2 <= t.crown_radius**2
+            if not m.any():
+                c, r = self.terrain.index_of(t.x, t.y)
+                m[r, c] = True
+            top[m] = np.maximum(top[m], t.height)
+        return top, top > 0
+
+    def surface(self, include_trees: bool = False) -> Grid:
+        """Terrain plus building heights (and tree crowns if asked) as a :class:`Grid`."""
+        h, _ = self.building_rasters()
+        if include_trees and self.trees:
+            t, _ = self.tree_raster()
+            h = np.maximum(h, t)
+        return self.terrain.copy(np.nan_to_num(self.terrain.values) + h)
+
+    # ------------------------------------------------------------ export
+    def to_manifest(self) -> dict[str, Any]:
+        g = self.terrain
+        width = g.cols * g.cell
+        height = g.rows * g.cell
+        valid = np.isfinite(g.values)
+        terrain: dict[str, Any] = {
+            "file": "terrain.bin",
+            "source": "zero2geospatial3d-sdk",
+            "min": float(np.nanmin(g.values)) if valid.any() else 0.0,
+            "max": float(np.nanmax(g.values)) if valid.any() else 0.0,
+        }
+        if not valid.all():
+            terrain["mask"] = "terrain_mask.bin"
+            terrain["nodataCells"] = int((~valid).sum())
+        return {
+            "format": SCENE_FORMAT,
+            "version": SCENE_VERSION,
+            "title": self.title,
+            "created": _dt.datetime.now().isoformat(timespec="seconds"),
+            "crs": self.crs,
+            "origin": list(self.origin),
+            "extent": {
+                "minX": -width / 2,
+                "minY": -height / 2,
+                "maxX": width / 2,
+                "maxY": height / 2,
+            },
+            "geo": {"lat": self.lat, "lon": self.lon, "utcOffset": self.utc_offset},
+            "grid": {
+                "cols": g.cols,
+                "rows": g.rows,
+                "cell": g.cell,
+                "minX": g.min_x,
+                "maxY": g.max_y,
+            },
+            "terrain": terrain,
+            "layers": [],
+            "rasters": [],
+            "weather": None,
+            "warnings": [],
+        }
+
+    def export(self, out_dir: str, epw_path: Optional[str] = None) -> dict[str, Any]:
+        """Write the viewer package to ``out_dir`` and return the manifest."""
+        os.makedirs(out_dir, exist_ok=True)
+        manifest = self.to_manifest()
+        values = self.terrain.values
+        valid = np.isfinite(values)
+        filled = np.where(valid, values, np.nanmin(values) if valid.any() else 0.0)
+        filled.astype("<f4").tofile(os.path.join(out_dir, "terrain.bin"))
+        if not valid.all():
+            valid.astype(np.uint8).tofile(os.path.join(out_dir, "terrain_mask.bin"))
+
+        def write_layer(
+            index: int,
+            name: str,
+            kind: str,
+            role: str,
+            feats: list[dict[str, Any]],
+            default_h: float,
+            color: str,
+        ) -> None:
+            fname = f"L{index}.geojson"
+            with open(os.path.join(out_dir, fname), "w", encoding="utf-8") as handle:
+                json.dump(
+                    {"type": "FeatureCollection", "features": feats}, handle, separators=(",", ":")
+                )
+            with open(os.path.join(out_dir, f"L{index}.ids.json"), "w", encoding="utf-8") as handle:
+                json.dump(list(range(1, len(feats) + 1)), handle)
+            field_meta = []
+            seen = set()
+            for f in feats:
+                for k, v in f["properties"].items():
+                    if k.startswith("_") or k in seen:
+                        continue
+                    seen.add(k)
+                    num = isinstance(v, (int, float)) and not isinstance(v, bool)
+                    field_meta.append({"name": k, "type": "number" if num else "text"})
+            manifest["layers"].append(
+                {
+                    "id": f"L{index}",
+                    "name": name,
+                    "file": fname,
+                    "kind": kind,
+                    "role": role,
+                    "scenario": "both",
+                    "count": len(feats),
+                    "fields": field_meta,
+                    "heightExpr": "",
+                    "baseExpr": "",
+                    "defaultHeight": default_h,
+                    "colorMode": "single",
+                    "defaultColor": color,
+                    "sourceLayerId": f"sdk_{role}",
+                    "sourceCrs": self.crs.get("authid", ""),
+                }
+            )
+
+        index = 0
+        if self.buildings:
+            feats = []
+            for fid, b in enumerate(self.buildings, start=1):
+                props: dict[str, Any] = {
+                    "_id": fid,
+                    "_h": round(max(0.0, b.height), 3),
+                    "_c": b.color,
+                }
+                if b.base > 0:
+                    props["_b"] = round(b.base, 3)
+                if b.group:
+                    props["_grp"] = str(b.group)
+                props.update(_plain_props(b.properties))
+                coords = [
+                    [
+                        [[round(float(p[0]), 2), round(float(p[1]), 2)] for p in ring]
+                        for ring in poly
+                    ]
+                    for poly in b.polygons
+                ]
+                feats.append(
+                    {
+                        "type": "Feature",
+                        "properties": props,
+                        "geometry": {"type": "MultiPolygon", "coordinates": coords},
+                    }
+                )
+            index += 1
+            write_layer(index, "Buildings", "polygon", "buildings", feats, 9.0, "#d9d4cc")
+        if self.trees:
+            feats = []
+            for fid, t in enumerate(self.trees, start=1):
+                props = {"_id": fid, "_h": round(t.height, 3), "_c": "#4f8a3c"}
+                props.update(_plain_props(t.properties))
+                feats.append(
+                    {
+                        "type": "Feature",
+                        "properties": props,
+                        "geometry": {"type": "MultiPoint", "coordinates": [[t.x, t.y]]},
+                    }
+                )
+            index += 1
+            write_layer(index, "Trees", "point", "trees", feats, 8.0, "#4f8a3c")
+        if self.roads:
+            feats = []
+            for fid, road in enumerate(self.roads, start=1):
+                props = {
+                    "_id": fid,
+                    "_h": 1.0,
+                    "_c": "#ff4f2e",
+                    "traffic_weight": road.traffic_weight,
+                }
+                props.update(_plain_props(road.properties))
+                line = [[[float(p[0]), float(p[1])] for p in road.coords]]
+                feats.append(
+                    {
+                        "type": "Feature",
+                        "properties": props,
+                        "geometry": {"type": "MultiLineString", "coordinates": line},
+                    }
+                )
+            index += 1
+            write_layer(index, "Roads", "line", "roads", feats, 1.0, "#ff4f2e")
+        if epw_path:
+            import shutil
+
+            shutil.copyfile(epw_path, os.path.join(out_dir, "weather.epw"))
+            manifest["weather"] = {"file": "weather.epw", "source": os.path.basename(epw_path)}
+        with open(os.path.join(out_dir, "scene.json"), "w", encoding="utf-8") as handle:
+            json.dump(manifest, handle, ensure_ascii=False, indent=1)
+        return manifest
+
+
+def _plain_props(props: dict[str, Any]) -> dict[str, Any]:
+    out: dict[str, Any] = {}
+    for k, v in props.items():
+        if isinstance(v, (bool, int)) or v is None:
+            out[k] = v
+        elif isinstance(v, float):
+            out[k] = v if math.isfinite(v) else None
+        else:
+            out[k] = str(v)[:120]
+    return out
+
+
+def load_scene(scene_dir: str) -> Scene:
+    """Read a scene package written by the plugin or by :meth:`Scene.export`."""
+    with open(os.path.join(scene_dir, "scene.json"), encoding="utf-8") as handle:
+        manifest = json.load(handle)
+    if manifest.get("format") != SCENE_FORMAT:
+        raise ValueError(f"{scene_dir} is not a {SCENE_FORMAT} package")
+    g = manifest["grid"]
+    cols, rows = int(g["cols"]), int(g["rows"])
+    values: np.ndarray = np.zeros((rows, cols))
+    terrain = manifest.get("terrain")
+    if terrain:
+        values = np.fromfile(os.path.join(scene_dir, terrain["file"]), dtype="<f4").astype(float)
+        values = values.reshape(rows, cols)
+        if terrain.get("mask"):
+            mask = np.fromfile(os.path.join(scene_dir, terrain["mask"]), dtype=np.uint8)
+            values[mask.reshape(rows, cols) == 0] = np.nan
+    grid = Grid(values, float(g["cell"]), float(g["minX"]), float(g["maxY"]))
+    buildings: list[Building] = []
+    trees: list[Tree] = []
+    roads: list[Road] = []
+    for layer in manifest.get("layers", []):
+        if layer.get("scenario", "both") == "B":
+            continue
+        with open(os.path.join(scene_dir, layer["file"]), encoding="utf-8") as handle:
+            fc = json.load(handle)
+        role = layer.get("role")
+        default_h = float(layer.get("defaultHeight") or 0.0)
+        for f in fc.get("features", []):
+            p = f.get("properties", {})
+            h = p.get("_h")
+            h = default_h if h is None else float(h)
+            geom = f.get("geometry") or {}
+            coords = geom.get("coordinates") or []
+            user = {k: v for k, v in p.items() if not k.startswith("_")}
+            if role == "buildings" and geom.get("type") in ("MultiPolygon", "Polygon"):
+                polys = coords if geom["type"] == "MultiPolygon" else [coords]
+                buildings.append(
+                    Building(
+                        polys,
+                        h,
+                        float(p.get("_b") or 0.0),
+                        color=p.get("_c", "#d9d4cc"),
+                        properties=user,
+                        group=p.get("_grp"),
+                    )
+                )
+            elif role == "trees":
+                pts = coords if geom.get("type") == "MultiPoint" else [coords]
+                for pt in pts:
+                    trees.append(
+                        Tree(
+                            float(pt[0]),
+                            float(pt[1]),
+                            h,
+                            max(1.5, 0.35 * h),
+                            0.3 * h,
+                            properties=user,
+                        )
+                    )
+            elif role == "roads":
+                lines = coords if geom.get("type") == "MultiLineString" else [coords]
+                w = float(user.get("traffic_weight", h or 1.0) or 1.0)
+                for line in lines:
+                    roads.append(Road(line, w, user))
+    geo_info = manifest.get("geo") or {}
+    return Scene(
+        grid,
+        buildings,
+        trees,
+        roads,
+        lat=float(geo_info.get("lat", 0.0)),
+        lon=float(geo_info.get("lon", 0.0)),
+        utc_offset=geo_info.get("utcOffset"),
+        title=manifest.get("title", "Scene"),
+        crs=manifest.get("crs"),
+        origin=tuple(manifest.get("origin", (0.0, 0.0))),
+    )