h3_utils.py

import numpy as np
import h3
import einops


def get_h3_grid_at_resolution(resolution: int) -> list[str]:
    """Get h3 cells (their h3 index) at a given resolution.
    Each cell appears once
    - resolution 0:       122 cells, every   ~20 degrees
    - resolution 1:       842 cells, every  ~7.5 degrees
    - resolution 2:     5,882 cells, every    ~3 degrees
    - resolution 3:    41,162 cells, every    ~1 degrees
    - resolution 4:   288,122 cells, every  ~0.4 degrees
    - resolution 5: 2,016,842 cells, every ~0.17 degrees
    c.f. https://h3geo.org/docs/core-library/restable
    """
    res0 = h3.get_res0_indexes()
    if resolution == 0:
        h = list(res0)
    else:
        h = [h3.h3_to_children(idx, resolution) for idx in res0]
        h = [item for sublist in h for item in sublist]  # flatten
    return h


def get_xyz_grid_at_resolution(resolution: int) -> np.ndarray:
    h3_grid = get_h3_grid_at_resolution(resolution)
    return np.array([h3_to_xyz(h) for h in h3_grid]).reshape((-1, 3))


def latlon_to_xyz(latlon: np.ndarray):
    latlon = np.array(latlon, dtype=np.float32)

    # Convert latitude and longitude from degrees to radians
    lat_radians = np.atleast_1d(np.deg2rad(latlon[..., 0]))
    lon_radians = np.atleast_1d(np.deg2rad(latlon[..., 1]))

    # Convert geographic coordinates to cartesian
    x = np.cos(lat_radians) * np.cos(lon_radians)
    y = np.cos(lat_radians) * np.sin(lon_radians)
    z = np.sin(lat_radians)
    return einops.rearrange([x, y, z], 'xyz ... -> ... xyz')


def xyz_to_latlon(xyz: np.ndarray):
    # Convert cartesian coordinates to geographic
    lat_radians = np.arcsin(xyz[..., 2])
    lon_radians = np.arctan2(xyz[..., 1], xyz[..., 0])

    # Convert latitude and longitude from radians to degrees
    latlon = einops.rearrange([lat_radians, lon_radians], 'latlon ... -> ... latlon')
    return np.rad2deg(latlon)


def h3_to_xyz(h: str):
    return latlon_to_xyz(np.array(h3.h3_to_geo(h)))


def plane_normal_to_rotation_matrix(plane_normal: np.ndarray):
    plane_normal = np.asarray(plane_normal)
    # generate x and y in the plane perpendicular to plane normal
    random_vector = np.random.uniform(size=(3, ))
    random_vector /= np.linalg.norm(random_vector)
    while np.any(np.dot(plane_normal, random_vector) == 1):
        random_vector = np.random.uniform(size=(3, ))
        random_vector /= np.linalg.norm(random_vector)
    y = np.cross(plane_normal, random_vector)
    y = y / np.linalg.norm(y, axis=-1, keepdims=True)
    x = np.cross(y, plane_normal)
    x = x / np.linalg.norm(x, axis=-1, keepdims=True)

    # construct rotation matrix
    rotation_matrix = np.zeros((len(x), 3, 3))
    rotation_matrix[:, :, 0] = x
    rotation_matrix[:, :, 1] = y
    rotation_matrix[:, :, 2] = plane_normal
    return rotation_matrix