dem-generate-rix-gdal.py

from pathlib import Path

import numpy as np
from numpy.typing import NDArray
from osgeo import gdal, ogr, osr
from pyproj import Transformer, CRS

_SECTOR_COUNT = 12
_SECTOR_SUBDIVISION_COUNT = 6
_SECTOR_RAY_LENGTH = 3500


def create_sector_angles() -> NDArray[np.floating]:
    sectors = np.linspace(0, 360, num=_SECTOR_COUNT * _SECTOR_SUBDIVISION_COUNT, endpoint=False)
    sectors = sectors - (360 / _SECTOR_COUNT) / 2  # The first sector CENTER is at 0°
    sectors = sectors.reshape((_SECTOR_COUNT, _SECTOR_SUBDIVISION_COUNT))
    return sectors % 360


def gdal_create_grid(ne: tuple[float, float], sw: tuple[float, float], crs: CRS, resolution: int):
    crs_to_proxy = Transformer.from_crs(crs, "EPSG:3035")  # to metric
    proxy_to_crs = Transformer.from_crs("EPSG:3035", crs)  # from metric

    ne = crs_to_proxy.transform(*ne)
    sw = crs_to_proxy.transform(*sw)

    x = np.arange(sw[0], ne[0], resolution)
    y = np.arange(sw[1], ne[1], resolution)

    x, y = np.meshgrid(x, y)

    xt, yt = proxy_to_crs.transform(x.flatten(), y.flatten())
    x, y = xt.reshape(x.shape), yt.reshape(y.shape)

    return np.vstack((x.ravel(), y.ravel())).T


def gdal_create_sector_lines(points: NDArray[np.floating], angles: NDArray[np.floating], crs: CRS):
    geod = crs.get_geod()

    crs_to_proxy = Transformer.from_crs(crs, "EPSG:4326")  # to metric
    proxy_to_crs = Transformer.from_crs("EPSG:4326", crs)  # from metric

    x, y = points.T
    x, y = crs_to_proxy.transform(x, y)

    xe, ye, _ = geod.fwd(
        np.tile(x, _SECTOR_SUBDIVISION_COUNT),
        np.tile(y, _SECTOR_SUBDIVISION_COUNT),
        np.repeat(angles, len(points)),
        np.repeat(_SECTOR_RAY_LENGTH, _SECTOR_SUBDIVISION_COUNT * len(points))
    )

    xe, ye = proxy_to_crs.transform(xe, ye)

    return xe.reshape((len(angles), len(points))), ye.reshape((len(angles), len(points)))


def gdal_create_sectors(points: NDArray[np.floating], crs: CRS):
    driver = ogr.GetDriverByName("MEMORY")
    source = driver.CreateDataSource("Sectors.shp")

    srs = osr.SpatialReference()
    srs.ImportFromWkt(crs.to_wkt())

    layers = []
    field_point_id = ogr.FieldDefn("PointId", ogr.OFTInteger)
    for i, sector in enumerate(create_sector_angles()):
        layer = source.CreateLayer(f"RaysSector{i}", srs, ogr.wkbLineString)
        layer.CreateField(field_point_id)

        line_end_x, line_end_y = gdal_create_sector_lines(points, sector, crs)
        for angle in range(_SECTOR_SUBDIVISION_COUNT):
            for j, ((px, py), lx, ly) in enumerate(zip(points, line_end_x[angle], line_end_y[angle])):
                line = ogr.Geometry(ogr.wkbLineString)
                line.AddPoint(px, py)
                line.AddPoint(lx, ly)

                feat = ogr.Feature(layer.GetLayerDefn())
                feat.SetGeometry(line)
                feat.SetField('PointId', j)

                layer.CreateFeature(feat)

        layers.append(layer)

    return source, layers


def gdal_main(
        raster_file=Path(r"hamburg10.tif"),
        contour_file=Path(r"hamburg10.shp"),
) -> None:
    raster = gdal.Open(str(raster_file), gdal.GA_ReadOnly)
    contours_ds = ogr.Open(str(contour_file))
    contours = contours_ds.GetLayer()

    crs = CRS.from_wkt(contours.GetSpatialRef().ExportToWkt())
    extent = contours.GetExtent()
    points = gdal_create_grid(
        ne=(extent[1], extent[3]),
        sw=(extent[0], extent[2]),
        crs=crs,
        resolution=1000
    )

    layer_source, layers = gdal_create_sectors(points, crs)

    driver = ogr.GetDriverByName("ESRI Shapefile")
    source = driver.CreateDataSource("Intersections.shp")
    srs = osr.SpatialReference()
    srs.ImportFromWkt(crs.to_wkt())

    for i, layer in enumerate(layers):
        intersection = source.CreateLayer(f"RayIntersection{i}", srs, ogr.wkbMultiPoint)
        contours.Intersection(layer, intersection, [
            "PROMOTE_TO_MULTI=YES",
            "KEEP_LOWER_DIMENSION_GEOMETRIES=YES"
        ], callback=lambda a, b, c: print(a, b, c))
        print(len(intersection))
        break


if __name__ == "__main__":
    gdal_main()