aaliddell / s2cell Goto Github PK

import math

from s2cell import s2cell  # type: ignore
from s2cell.s2cell import (  # type: ignore
    _S2_INVERT_MASK,
    _S2_LOOKUP_BITS,
    _S2_MAX_LEVEL,
    _S2_MAX_SIZE,
    _S2_POS_BITS,
    _S2_SWAP_MASK,
    _s2_face_uv_to_xyz,
    _s2_st_to_ij,
    cell_id_to_level,
)


def get_edge_neighbor_cell_ids(cellid: int) -> list[int]:
    level = cell_id_to_level(cellid)
    size = _get_size_ij(level)
    face, i, j = _to_face_ij_orientation(cellid)
    return [
        _from_face_ij_same(face, i, j - size, j - size >= 0, level),
        _from_face_ij_same(face, i + size, j, i + size < _S2_MAX_SIZE, level),
        _from_face_ij_same(face, i, j + size, j + size < _S2_MAX_SIZE, level),
        _from_face_ij_same(face, i - size, j, i - size >= 0, level),
    ]


def get_vertex_neighbor_cell_ids(cellid: int) -> list[int]:
    level = cell_id_to_level(cellid)
    size = _get_size_ij(level)
    face, i, j = _to_face_ij_orientation(cellid)
    neighbors = []
    for di in (-size, size):
        for dj in (-size, size):
            same = (0 <= (i + di) < _S2_MAX_SIZE) and (0 <= (j + dj) < _S2_MAX_SIZE)
            neighbors.append(_from_face_ij_same(face, i + di, j + dj, same, level))
    return neighbors


def _get_size_ij(level: int) -> int:
    return 1 << (_S2_MAX_LEVEL - level)


def _to_face_ij_orientation(cell_id: int) -> tuple[int, int, int]:
    s2cell._s2_init_lookups()

    # https://github.com/aaliddell/s2cell/blob/main/s2cell/s2cell.py#L546-L607
    face = cell_id >> _S2_POS_BITS
    bits = face & _S2_SWAP_MASK  # ppppppppoo. Initially set by by face
    lookup_mask = (1 << _S2_LOOKUP_BITS) - 1  # Mask of 4 one bits: 0b1111
    i = 0
    j = 0
    for k in range(7, -1, -1):
        # Pull out 8 bits of cell ID, except in first loop where we pull out only 4
        n_bits = (_S2_MAX_LEVEL - 7 * _S2_LOOKUP_BITS) if k == 7 else _S2_LOOKUP_BITS
        extract_mask = (1 << (2 * n_bits)) - 1  # 8 (or 4) one bits
        bits += ((cell_id >> (k * 2 * _S2_LOOKUP_BITS + 1)) & extract_mask) << 2

        # Map bits from ppppppppoo to iiiijjjjoo using lookup table
        bits = s2cell._S2_LOOKUP_IJ[bits]

        # Extract I and J bits
        offset = k * _S2_LOOKUP_BITS
        i += (bits >> (_S2_LOOKUP_BITS + 2)) << offset  # Don't need lookup mask here
        j += ((bits >> 2) & lookup_mask) << offset

        # Remove I and J bits, leaving just new swap and invert bits for the next round
        bits &= _S2_SWAP_MASK | _S2_INVERT_MASK  # Mask: 0b11
    return face, i, j


def _get_face(p: tuple[float, float, float]) -> int:
    # https://github.com/aaliddell/s2cell/blob/main/s2cell/s2cell.py#L389-L391
    face = max(enumerate(p), key=lambda p: abs(p[1]))[0]  # Largest absolute component
    if p[face] < 0.0:
        face += 3
    return face


def _xyz_to_face_uv(p: tuple[float, float, float]) -> tuple[int, float, float]:
    face = _get_face(p)
    # https://github.com/aaliddell/s2cell/blob/main/s2cell/s2cell.py#L393-L423
    uv = (  # pylint: disable=invalid-name
        p[1 - ((face + 1) >> 1)] / p[face % 3],  # U
        p[2 - (face >> 1)] / p[face % 3],  # V
    )
    if face in (1, 2, 5):
        uv = (-uv[0], uv[1])  # Negate U  # pylint: disable=invalid-name
    if face in (2, 4, 5):
        uv = (uv[0], -uv[1])  # Negate V  # pylint: disable=invalid-name
    return face, uv[0], uv[1]
    # if face == 0:
    #     return face, p[1] / p[0], p[2] / p[0]
    # elif face == 1:
    #     return face, -p[0] / p[1], p[2] / p[1]
    # elif face == 2:
    #     return face, -p[0] / p[2], -p[1] / p[2]
    # elif face == 3:
    #     return face, p[2] / p[0], p[1] / p[0]
    # elif face == 4:
    #     return face, p[2] / p[1], -p[0] / p[1]
    # else:
    #     return face, -p[1] / p[2], -p[0] / p[2]


def _from_face_ij_same(face: int, i: int, j: int, same: bool, level: int) -> int:
    if same:
        return _from_face_ij(face, i, j, level)
    else:
        return _from_face_ij_wrap(face, i, j, level)


def _from_face_ij_wrap(face: int, i: int, j: int, level: int) -> int:
    # https://github.com/google/s2geometry/blob/c59d0ca01ae3976db7f8abdc83fcc871a3a95186/src/s2/s2cell_id.cc#L446-L477
    i = max(-1, min(_S2_MAX_SIZE, i))
    j = max(-1, max(_S2_MAX_SIZE, j))

    scale = 1.0 / _S2_MAX_SIZE
    limit = math.nextafter(1, 2)
    u = max(-limit, min(limit, scale * ((i << 1) + 1 - _S2_MAX_SIZE)))
    v = max(-limit, min(limit, scale * ((j << 1) + 1 - _S2_MAX_SIZE)))

    face, u, v = _xyz_to_face_uv(_s2_face_uv_to_xyz(face, (u, v)))
    return _from_face_ij(
        face, _s2_st_to_ij(0.5 * (u + 1)), _s2_st_to_ij(0.5 * (v + 1)), level
    )


def _from_face_ij(face: int, i: int, j: int, level: int) -> int:
    s2cell._s2_init_lookups()

    # https://github.com/aaliddell/s2cell/blob/main/s2cell/s2cell.py#L433-L484
    bits = face & _S2_SWAP_MASK  # iiiijjjjoo. Initially set by by face
    cell_id = face << (_S2_POS_BITS - 1)  # Insert face at second most signficant bits
    lookup_mask = (1 << int(_S2_LOOKUP_BITS)) - 1  # Mask of 4 one bits: 0b1111
    required_steps = math.ceil((level + 2) / 4) if level > 0 else 0
    for k in range(7, 7 - required_steps, -1):
        # Grab 4 bits of each of I and J
        offset = k * _S2_LOOKUP_BITS
        bits += ((i >> offset) & lookup_mask) << (_S2_LOOKUP_BITS + 2)
        bits += ((j >> offset) & lookup_mask) << 2

        # Map bits from iiiijjjjoo to ppppppppoo using lookup table
        bits = s2cell._S2_LOOKUP_POS[bits]

        # Insert position bits into cell ID
        cell_id |= (bits >> 2) << (k * 2 * _S2_LOOKUP_BITS)

        # Remove position bits, leaving just new swap and invert bits for the next round
        bits &= _S2_SWAP_MASK | _S2_INVERT_MASK  # Mask: 0b11

    # Left shift and add trailing bit
    # The trailing bit addition is disabled, as we are overwriting this below in the truncation
    # anyway. This line is kept as an example of the full method for S2 cell ID creation as is done
    # in the standard library versions.
    cell_id = cell_id << 1  # + 1

    # Truncate to desired level
    # This is done by finding the mask of the trailing 1 bit for the specified level, then zeroing
    # out all bits less significant than this, then finally setting the trailing 1 bit. This is
    # still necessary to do even after a reduced number of steps `required_steps` above, since each
    # step contains multiple levels that may need partial overwrite. Additionally, we need to add
    # the trailing 1 bit, which is not yet set above.
    least_significant_bit_mask = 1 << (2 * (_S2_MAX_LEVEL - level))
    cell_id = (cell_id & -least_significant_bit_mask) | least_significant_bit_mask

    return cell_id

if __name__ == "__main__":
    # lat, lng = 0, 0
    # lat, lng = -7.8556806, -34.8624166
    lat, lng = 36.7066928, -45.0000000
    level = 5
    from s2cell import lat_lon_to_cell_id, cell_id_to_token
    cell_id = lat_lon_to_cell_id(lat, lng, level)
    print(cell_id_to_token(cell_id))
    for neighbor in get_edge_neighbor_cell_ids(cell_id):
        print("  ", cell_id_to_token(neighbor))
    for neighbor in get_vertex_neighbor_cell_ids(cell_id):
        print("  ", cell_id_to_token(neighbor))