Setup cython instead of numba

.github/workflows/deploy.yaml

@@ -29,9 +29,9 @@ jobs:
             trusted-public-keys = cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY=
             substituters = https://cache.nixos.org/
 
-      - name: Install dependencies
+      - name: Install dependencies and compile cython
         run: |
-          nix-shell --run "true"
+          nix-shell --run "make"
 
       - name: Build
         run: |

.gitignore

@@ -199,3 +199,6 @@ pyrightconfig.json
 result
 data/*
 cert/*
+
+cython_lib/*.c
+cython_lib/*.html

Makefile

@@ -1,6 +1,9 @@
-.PHONY: build version dev-start dev-stop dev-logs
+.PHONY: setup docker version dev-start dev-stop dev-logs
 
-build:
+setup:
+	python setup.py build_ext --build-lib cython_lib
+
+docker:
 	docker load < $$(nix-build --no-out-link)
 
 version:

Pipfile

@@ -9,6 +9,7 @@ asyncache = "*"
 authlib = "*"
 brotlipy = "*"
 cachetools = "*"
+cython = "*"
 dacite = "*"
 fastapi = "*"
 feedgen = "*"
@@ -18,7 +19,6 @@ jinja2 = "*"
 mapbox-vector-tile = "*"
 motor = "*"
 networkx = "*"
-numba = "*"
 numpy = "*"
 orjson = "*"
 pillow = "*"

api/v1/tile.py

@@ -1,4 +1,3 @@
-import math
 from collections.abc import Sequence
 from itertools import chain
 from typing import Annotated
@@ -8,7 +7,6 @@
 import numpy as np
 from anyio.streams.memory import MemoryObjectSendStream
 from fastapi import APIRouter, Path, Response
-from numba import njit
 from shapely.ops import transform
 
 from config import (
@@ -22,33 +20,18 @@
     TILE_MAX_Z,
     TILE_MIN_Z,
 )
+from cython_lib.geo_utils import tile_to_bbox
 from middlewares.cache_middleware import make_cache_control
 from models.aed import AED
 from models.bbox import BBox
 from models.country import Country
-from models.lonlat import LonLat
 from states.aed_state import AEDState, AEDStateDep
 from states.country_state import CountryState, CountryStateDep
 from utils import abbreviate, print_run_time
 
 router = APIRouter()
 
 
-@njit(fastmath=True)
-def _tile_to_lonlat(z: int, x: int, y: int) -> tuple[float, float]:
-    n = 2.0**z
-    lon_deg = x / n * 360.0 - 180.0
-    lat_rad = math.atan(math.sinh(math.pi * (1 - 2 * y / n)))
-    lat_deg = math.degrees(lat_rad)
-    return lon_deg, lat_deg
-
-
-def _tile_to_bbox(z: int, x: int, y: int) -> BBox:
-    p1_lon, p1_lat = _tile_to_lonlat(z, x, y)
-    p2_lon, p2_lat = _tile_to_lonlat(z, x + 1, y + 1)
-    return BBox(LonLat(p1_lon, p2_lat), LonLat(p2_lon, p1_lat))
-
-
 async def _count_aed_in_country(country: Country, aed_state: AEDState, send_stream: MemoryObjectSendStream) -> None:
     count = await aed_state.count_aeds_by_country_code(country.code)
     await send_stream.send((country, count))
@@ -174,7 +157,7 @@ async def get_tile(
     country_state: CountryStateDep,
     aed_state: AEDStateDep,
 ):
-    bbox = _tile_to_bbox(z, x, y)
+    bbox = tile_to_bbox(z, x, y)
     assert bbox.p1.lon <= bbox.p2.lon, f'{bbox.p1.lon=} <= {bbox.p2.lon=}'
     assert bbox.p1.lat <= bbox.p2.lat, f'{bbox.p1.lat=} <= {bbox.p2.lat=}'
 

cython_lib/geo_utils.py

@@ -0,0 +1,33 @@
+import cython
+
+from models.bbox import BBox
+from models.lonlat import LonLat
+
+if cython.compiled:
+    from cython.cimports.libc.math import atan, pi, sinh
+
+    print(f'{__name__}: 🐇 compiled')
+else:
+    from math import atan, pi, sinh
+
+    print(f'{__name__}: 🐌 not compiled')
+
+
+@cython.cfunc
+def _degrees(radians: cython.double) -> cython.double:
+    return radians * 180 / pi
+
+
+@cython.cfunc
+def _tile_to_lonlat(z: cython.int, x: cython.int, y: cython.int) -> tuple[cython.double, cython.double]:
+    n: cython.double = 2**z
+    lon_deg: cython.double = x / n * 360.0 - 180.0
+    lat_rad: cython.double = atan(sinh(pi * (1 - 2 * y / n)))
+    lat_deg: cython.double = _degrees(lat_rad)
+    return lon_deg, lat_deg
+
+
+def tile_to_bbox(z: cython.int, x: cython.int, y: cython.int) -> BBox:
+    p1_lon, p1_lat = _tile_to_lonlat(z, x, y)
+    p2_lon, p2_lat = _tile_to_lonlat(z, x + 1, y + 1)
+    return BBox(LonLat(p1_lon, p2_lat), LonLat(p2_lon, p1_lat))

default.nix

@@ -35,11 +35,11 @@ dockerTools.buildLayeredImage {
     cp "${./.}"/LICENSE .
     cp "${./.}"/Makefile .
     cp "${./.}"/*.py .
-    mkdir -p api/v1 middlewares models states
-    cp "${./api/v1}"/*.py api/v1
-    cp "${./middlewares}"/*.py middlewares
-    cp "${./models}"/*.py models
-    cp "${./states}"/*.py states
+    cp -r "${./.}"/api .
+    cp -r "${./.}"/cython_lib .
+    cp -r "${./.}"/middlewares .
+    cp -r "${./.}"/models .
+    cp -r "${./.}"/states .
     ${shell.shellHook}
   '';
 

setup.py

@@ -0,0 +1,31 @@
+import Cython.Compiler.Options as Options
+from Cython.Build import cythonize
+from setuptools import Extension, setup
+
+Options.docstrings = False
+Options.annotate = True
+
+setup(
+    ext_modules=cythonize(
+        [
+            Extension(
+                '*',
+                ['cython_lib/*.py'],
+                extra_compile_args=[
+                    '-march=x86-64',
+                    '-ffast-math',
+                    '-fopenmp',
+                    '-flto=auto',
+                ],
+                extra_link_args=[
+                    '-fopenmp',
+                    '-flto=auto',
+                ],
+            ),
+        ],
+        compiler_directives={
+            # https://cython.readthedocs.io/en/latest/src/userguide/source_files_and_compilation.html#compiler-directives
+            'language_level': 3,
+        },
+    ),
+)

utils.py

@@ -5,12 +5,10 @@
 from contextlib import contextmanager
 from dataclasses import asdict
 from datetime import timedelta
-from math import atan2, cos, pi, radians, sin, sqrt
 from typing import Any
 
 import anyio
 import httpx
-from numba import njit
 from shapely.geometry import mapping
 
 from config import USER_AGENT
@@ -90,47 +88,3 @@ def as_dict(data) -> dict:
             d[k] = mapping(v)
 
     return d
-
-
-EARTH_RADIUS = 6371000
-CIRCUMFERENCE = 2 * pi * EARTH_RADIUS
-
-
-@njit(fastmath=True)
-def meters_to_lat(meters: float) -> float:
-    return meters / (CIRCUMFERENCE / 360)
-
-
-@njit(fastmath=True)
-def meters_to_lon(meters: float, lat: float) -> float:
-    return meters / ((CIRCUMFERENCE / 360) * cos(radians(lat)))
-
-
-@njit(fastmath=True)
-def lat_to_meters(lat: float) -> float:
-    return lat * (CIRCUMFERENCE / 360)
-
-
-@njit(fastmath=True)
-def lon_to_meters(lon: float, lat: float) -> float:
-    return lon * ((CIRCUMFERENCE / 360) * cos(radians(lat)))
-
-
-@njit(fastmath=True)
-def radians_tuple(p: tuple[float, float]) -> tuple[float, float]:
-    return (radians(p[0]), radians(p[1]))
-
-
-@njit(fastmath=True)
-def haversine_distance(p1: tuple[float, float], p2: tuple[float, float]) -> float:
-    p1_lat, p1_lon = radians_tuple(p1)
-    p2_lat, p2_lon = radians_tuple(p2)
-
-    dlat = p2_lat - p1_lat
-    dlon = p2_lon - p1_lon
-
-    a = sin(dlat / 2) ** 2 + cos(p1_lat) * cos(p2_lat) * sin(dlon / 2) ** 2
-    c = 2 * atan2(sqrt(a), sqrt(1 - a))
-
-    # distance in meters
-    return c * EARTH_RADIUS