diff --git a/app/broadcast_areas/__init__.py b/app/broadcast_areas/__init__.py index 230313694..84be1e62c 100644 --- a/app/broadcast_areas/__init__.py +++ b/app/broadcast_areas/__init__.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from notifications_utils.formatters import formatted_list from notifications_utils.polygons import Polygons from notifications_utils.serialised_model import SerialisedModelCollection +from rtreelib import Rect from werkzeug.utils import cached_property from .populations import CITY_OF_LONDON @@ -157,9 +158,11 @@ class CustomBroadcastArea(BaseBroadcastArea): def overlapping_areas(self): if not self.polygons: return [] - return broadcast_area_libraries.get_areas( - *rtree_index.intersection(self.polygons.bounds, objects='raw') - ) + return broadcast_area_libraries.get_areas([ + overlap.data for overlap in rtree_index.query( + Rect(*self.polygons.bounds) + ) + ]) @cached_property def count_of_phones(self): diff --git a/app/broadcast_areas/create-broadcast-areas-db.py b/app/broadcast_areas/create-broadcast-areas-db.py index a404e5312..c8ead3137 100755 --- a/app/broadcast_areas/create-broadcast-areas-db.py +++ b/app/broadcast_areas/create-broadcast-areas-db.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import csv +import pickle import sys from pathlib import Path @@ -15,10 +16,12 @@ from populations import ( SMARTPHONE_OWNERSHIP_BY_AGE_RANGE, estimate_number_of_smartphones_for_population, ) -from repo import BroadcastAreasRepository, rtree_index +from repo import BroadcastAreasRepository, rtree_index_path +from rtreelib import Rect, RTree source_files_path = Path(__file__).resolve().parent / 'source_files' point_counts = [] +rtree_index = RTree() def simplify_geometry(feature): @@ -227,7 +230,7 @@ def add_wards_local_authorities_and_counties(): def _add_electoral_wards(dataset_id): areas_to_add = [] - for index, feature in enumerate(geojson.loads(wd20_filepath.read_text())["features"]): + for feature in geojson.loads(wd20_filepath.read_text())["features"]: ward_code = feature["properties"]["wd20cd"] ward_name = feature["properties"]["wd20nm"] ward_id = "wd20-" + ward_code @@ -243,7 +246,7 @@ def _add_electoral_wards(dataset_id): ) if feature: - rtree_index.insert(index, Polygons(feature).bounds, obj=ward_id) + rtree_index.insert(ward_id, Rect(*Polygons(feature).bounds)) areas_to_add.append([ ward_id, ward_name, @@ -255,6 +258,7 @@ def _add_electoral_wards(dataset_id): except KeyError: print("Skipping", ward_code, ward_name) # noqa: T001 + rtree_index_path.open('wb').write(pickle.dumps(rtree_index)) repo.insert_broadcast_areas(areas_to_add, keep_old_polygons) @@ -340,6 +344,4 @@ print( # noqa: T001 'DONE\n' f' Processed {len(point_counts):,} polygons.\n' f' Highest point counts once simplifed: {most_detailed_polygons}\n' - f' RTree bounds: {rtree_index.bounds}\n' - f' Number of objects in Rtree: {rtree_index.get_size():,}\n' ) diff --git a/app/broadcast_areas/repo.py b/app/broadcast_areas/repo.py index 0733c8a6c..91014077a 100644 --- a/app/broadcast_areas/repo.py +++ b/app/broadcast_areas/repo.py @@ -1,14 +1,11 @@ import json import os +import pickle import sqlite3 from pathlib import Path -from rtree import index - -rtree_index = index.Rtree( - str((Path(__file__).parent / 'rtree').absolute()), - interleaved=True, -) +rtree_index_path = Path(__file__).parent / 'rtree.pickle' +rtree_index = pickle.loads(rtree_index_path.read_bytes()) class BroadcastAreasRepository(object): diff --git a/app/broadcast_areas/rtree.dat b/app/broadcast_areas/rtree.dat deleted file mode 100644 index fc5a04591..000000000 Binary files a/app/broadcast_areas/rtree.dat and /dev/null differ diff --git a/app/broadcast_areas/rtree.idx b/app/broadcast_areas/rtree.idx deleted file mode 100644 index 24f87c9cf..000000000 Binary files a/app/broadcast_areas/rtree.idx and /dev/null differ diff --git a/app/broadcast_areas/rtree.pickle b/app/broadcast_areas/rtree.pickle new file mode 100644 index 000000000..0eedc6b7e Binary files /dev/null and b/app/broadcast_areas/rtree.pickle differ diff --git a/requirements.in b/requirements.in index af3517dae..9e26a8c46 100644 --- a/requirements.in +++ b/requirements.in @@ -19,7 +19,7 @@ gunicorn==20.1.0 eventlet==0.30.2 notifications-python-client==6.0.2 Shapely==1.7.1 -Rtree==0.9.7 +rtreelib==0.2.0 # PaaS awscli-cwlogs>=1.4,<1.5 diff --git a/requirements.txt b/requirements.txt index 3bbb6acbe..37a0c32b5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -171,7 +171,7 @@ requests==2.25.1 # notifications-utils rsa==4.7.2 # via awscli -rtree==0.9.7 +rtreelib==0.2.0 # via -r requirements.in s3transfer==0.3.6 # via