uphill-rivers/processing/run.py

import os
import json
import pickle
import hashlib
import overpy
import sqlite3
import numpy as np

from river import River, RiverTooShort

SQLITE_DB_FILE = 'db.sqlite'


def get_world_bbox(db: 'DBClient') -> list[tuple[float, float, float, float]]:
    """ Segments the world into bboxes, to be used by the overpass api.
    """
    #for lat in np.arang(-90, 90, 0.1):
    #    for lon in np.arang(-180, 180, 0.1):
    for lat in np.arange(46, 47, 0.1):
        for lon in np.arange(6, 7, 0.1):
            """ Check if bbox was done in this run"""
            elat = lat + 0.1
            elon = lon + 0.1
            lat = round(lat, 2)
            lon = round(lon, 2)
            elat = round(elat, 2)
            elon = round(elon, 2)

            print(f'Generated BBOX {lat}, {lon}, {elat}, {elon}')
            if db.has_bbox_run_this_run(lat, lon, elat, elon):
                print('BBOX already done this run. Will continue')
                continue

            yield (lat, lon, elat, elon)

            print('Finished processing bbox. Will continue to next one')
            db.done_bbox(lat, lon, elat, elon)

def overpass_cache(query: str) -> overpy.Result:
    # Only used for testing purposes
    query_hash = hashlib.sha256(query.encode()).hexdigest()
    fn = f'cache/{query_hash}.pickle'

    if not os.path.exists(fn):
        with open(fn, 'wb') as f:
            print("Will query overpass")

            rslt = overpy.Overpass(url='https://overpass.kumi.systems/api/interpreter').query(query)
            pickle.dump(rslt, f)
    else:
        print("Will load from cache")

    with open(fn, 'rb') as f:
        return pickle.load(f)

def get_rivers_of_bbox(bbox) -> list[River]:
    bstr = f'{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]}'

    result = overpass_cache(f"""
    [timeout:25];
    (

    way["waterway"="river"]({bstr});
    relation["waterway"="river"]({bstr});
    );
    out body;
    >;
    out skel qt;
        """)

    for way in result.ways:
        wayid = way.id
        r = River(wayid, way)
        yield r

def process_river(r: River):
    print(f'\nProcessing river {r.wayid}')

    results_file = f'results/{r.wayid}.json'
    try:
        r.check_slope()
        if not r.is_correct_slope:
            print(f"River is not downhill. Confidence: {r.correct_slope_confidence}")
            print(r.osm_link())
            with open(results_file, 'w') as f:
                js = r.to_dict()
                f.write(json.dumps(js, indent=1))
    except RiverTooShort:
        print("River is too short")
        print(r.osm_link())


class DBClient:
    runid: int
    def __init__(self, runid):
        self.runid = runid
        self.conn = self._get_db()

    def _get_db(self) -> sqlite3.Connection:
        if not os.path.exists(SQLITE_DB_FILE):
            conn = sqlite3.connect(SQLITE_DB_FILE)
            c = conn.cursor()
            c.execute('CREATE TABLE worldexport (slat REAL, slon REAL, elat REAL, elon REAL, last_run INTEGER, PRIMARY KEY (slat, slon, elat, elon))')
            c.execute('CREATE TABLE rivers (wayid INTEGER PRIMARY KEY, name TEXT, length REAL, start_ele REAL, end_ele REAL, slope_correct INTEGER, slope_confidence REAL, centroid_lat REAL, centroid_lon REAL, last_run INTEGER)')
            conn.commit()
            conn.close()

        return sqlite3.connect(SQLITE_DB_FILE)

    def has_bbox_run_this_run(self, slat: float, slon: float, elat: float, elon: float) -> int:
        c = self.conn.cursor()
        c.execute('SELECT last_run FROM worldexport WHERE slat = ? AND slon = ? AND elat = ? AND elon = ?', (slat, slon, elat, elon))
        results = c.fetchall()
        if results is None or results == []:
            return False
        return results[0][0] == self.runid

    def done_bbox(self, slat: float, slon: float, elat: float, elon: float):
        c = self.conn.cursor()
        c.execute('REPLACE INTO worldexport (slat, slon, elat, elon, last_run) VALUES (?, ?, ?, ?, ?)', (slat, slon, elat, elon, self.runid))
        self.conn.commit()

    def is_river_done(self, r: River) -> bool:
        c = self.conn.cursor()
        c.execute('SELECT last_run FROM rivers WHERE wayid = ?', (r.wayid,))
        results = c.fetchall()
        if results is None or results == []:
            return False
        return results[0][0] == self.runid

    def insert_river(self, r: River):
        c = self.conn.cursor()
        centroid = r.centroid
        clat = float(centroid[0])
        clon = float(centroid[1])
        c.execute('''
                  REPLACE INTO rivers
                  (wayid, name, length, start_ele, end_ele, slope_correct, slope_confidence, centroid_lat, centroid_lon, last_run)
                  VALUES
                  (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
                  (r.wayid, r.name, r.length, r.start_ele, r.end_ele, r.is_correct_slope, r.correct_slope_confidence, clat, clon, self.runid))
        self.conn.commit()


if __name__ == '__main__':
    last_run = 1
    db = DBClient(last_run)
    for bbox in get_world_bbox(db):
        for r in get_rivers_of_bbox(bbox):
            if not db.is_river_done(r):
                process_river(r)
                db.insert_river(r)