diff --git a/app/broadcast_areas/__init__.py b/app/broadcast_areas/__init__.py index caf94aee9..230313694 100644 --- a/app/broadcast_areas/__init__.py +++ b/app/broadcast_areas/__init__.py @@ -7,7 +7,7 @@ from notifications_utils.serialised_model import SerialisedModelCollection from werkzeug.utils import cached_property from .populations import CITY_OF_LONDON -from .repo import BroadcastAreasRepository +from .repo import BroadcastAreasRepository, rtree_index class SortableMixin: @@ -139,10 +139,6 @@ class BroadcastArea(BaseBroadcastArea, SortableMixin): class CustomBroadcastArea(BaseBroadcastArea): - # We don’t yet have a way to estimate the number of phones in a - # user-defined polygon - count_of_phones = 0 - def __init__(self, *, name, polygons=None): self.name = name self._polygons = polygons or [] @@ -157,6 +153,21 @@ class CustomBroadcastArea(BaseBroadcastArea): simple_polygons = polygons + @property + def overlapping_areas(self): + if not self.polygons: + return [] + return broadcast_area_libraries.get_areas( + *rtree_index.intersection(self.polygons.bounds, objects='raw') + ) + + @cached_property + def count_of_phones(self): + return sum( + area.polygons.ratio_of_intersection_with(self.polygons) * area.count_of_phones + for area in self.overlapping_areas + ) + class CustomBroadcastAreas(SerialisedModelCollection): model = CustomBroadcastArea diff --git a/app/broadcast_areas/create-broadcast-areas-db.py b/app/broadcast_areas/create-broadcast-areas-db.py index f613895dd..a404e5312 100755 --- a/app/broadcast_areas/create-broadcast-areas-db.py +++ b/app/broadcast_areas/create-broadcast-areas-db.py @@ -15,7 +15,7 @@ from populations import ( SMARTPHONE_OWNERSHIP_BY_AGE_RANGE, estimate_number_of_smartphones_for_population, ) -from repo import BroadcastAreasRepository +from repo import BroadcastAreasRepository, rtree_index source_files_path = Path(__file__).resolve().parent / 'source_files' point_counts = [] @@ -227,7 +227,7 @@ def add_wards_local_authorities_and_counties(): def _add_electoral_wards(dataset_id): areas_to_add = [] - for feature in geojson.loads(wd20_filepath.read_text())["features"]: + for index, feature in enumerate(geojson.loads(wd20_filepath.read_text())["features"]): ward_code = feature["properties"]["wd20cd"] ward_name = feature["properties"]["wd20nm"] ward_id = "wd20-" + ward_code @@ -242,6 +242,9 @@ def _add_electoral_wards(dataset_id): polygons_and_simplified_polygons(feature["geometry"]) ) + if feature: + rtree_index.insert(index, Polygons(feature).bounds, obj=ward_id) + areas_to_add.append([ ward_id, ward_name, dataset_id, la_id, @@ -337,4 +340,6 @@ 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 d6f4ebcf7..0733c8a6c 100644 --- a/app/broadcast_areas/repo.py +++ b/app/broadcast_areas/repo.py @@ -3,6 +3,13 @@ import os import sqlite3 from pathlib import Path +from rtree import index + +rtree_index = index.Rtree( + str((Path(__file__).parent / 'rtree').absolute()), + interleaved=True, +) + class BroadcastAreasRepository(object): def __init__(self): diff --git a/app/broadcast_areas/rtree.dat b/app/broadcast_areas/rtree.dat new file mode 100644 index 000000000..fc5a04591 Binary files /dev/null and b/app/broadcast_areas/rtree.dat differ diff --git a/app/broadcast_areas/rtree.idx b/app/broadcast_areas/rtree.idx new file mode 100644 index 000000000..24f87c9cf Binary files /dev/null and b/app/broadcast_areas/rtree.idx differ diff --git a/requirements.in b/requirements.in index 37d41edc9..af3517dae 100644 --- a/requirements.in +++ b/requirements.in @@ -19,12 +19,13 @@ gunicorn==20.1.0 eventlet==0.30.2 notifications-python-client==6.0.2 Shapely==1.7.1 +Rtree==0.9.7 # PaaS awscli-cwlogs>=1.4,<1.5 itsdangerous==1.1.0 -git+https://github.com/alphagov/notifications-utils.git@44.1.0#egg=notifications-utils==44.1.0 +git+https://github.com/alphagov/notifications-utils.git@44.2.0#egg=notifications-utils==44.2.0 git+https://github.com/alphagov/govuk-frontend-jinja.git@v0.5.8-alpha#egg=govuk-frontend-jinja==0.5.8-alpha # gds-metrics requires prometheseus 0.2.0, override that requirement as later versions bring significant performance gains diff --git a/requirements.txt b/requirements.txt index 137ae320c..3bbb6acbe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -108,7 +108,7 @@ mistune==0.8.4 # via notifications-utils notifications-python-client==6.0.2 # via -r requirements.in -git+https://github.com/alphagov/notifications-utils.git@44.1.0#egg=notifications-utils==44.1.0 +git+https://github.com/alphagov/notifications-utils.git@44.2.0#egg=notifications-utils==44.2.0 # via -r requirements.in openpyxl==3.0.7 # via pyexcel-xlsx @@ -171,6 +171,8 @@ requests==2.25.1 # notifications-utils rsa==4.7.2 # via awscli +rtree==0.9.7 + # via -r requirements.in s3transfer==0.3.6 # via # awscli diff --git a/tests/app/broadcast_areas/custom_polygons.py b/tests/app/broadcast_areas/custom_polygons.py new file mode 100644 index 000000000..6faea8675 --- /dev/null +++ b/tests/app/broadcast_areas/custom_polygons.py @@ -0,0 +1,15 @@ +BRISTOL = [ + [51.4371, -2.6216], + [51.4371, -2.5750], + [51.4668, -2.5750], + [51.4668, -2.6216], + [51.4371, -2.6216], +] + +SKYE = [ + [57.1004, -6.8280], + [57.1004, -5.7733], + [57.7334, -5.7733], + [57.7334, -6.8280], + [57.1004, -6.8280], +] diff --git a/tests/app/broadcast_areas/test_broadcast_area.py b/tests/app/broadcast_areas/test_broadcast_area.py index e612fd976..f25a44856 100644 --- a/tests/app/broadcast_areas/test_broadcast_area.py +++ b/tests/app/broadcast_areas/test_broadcast_area.py @@ -1,9 +1,11 @@ from math import isclose import pytest +from custom_polygons import BRISTOL, SKYE from app.broadcast_areas import ( BroadcastAreasRepository, + CustomBroadcastArea, broadcast_area_libraries, ) from app.broadcast_areas.populations import ( @@ -350,3 +352,50 @@ def test_estimated_bleed( broadcast_area_libraries.get_areas(area)[0].estimated_bleed_in_degrees, expected_bleed_in_degrees, ) + + +@pytest.mark.parametrize('polygon, expected_possible_overlaps, expected_count_of_phones', ( + ( + BRISTOL, + [ + 'Ashley', + 'Bedminster', + 'Central', + 'Clifton', + 'Clifton Down', + 'Cotham', + 'Hotwells and Harbourside', + 'Knowle', + 'Lawrence Hill', + 'Southville', + 'Stoke Bishop', + 'Windmill Hill', + ], + 73_496, + ), + ( + SKYE, + [ + 'Caol and Mallaig', + 'Eilean á Chèo', + 'Na Hearadh agus Ceann a Deas nan Loch', + 'Wester Ross, Strathpeffer and Lochalsh', + ], + 3_517, + ), +)) +def test_count_of_phones_for_custom_area( + polygon, + expected_possible_overlaps, + expected_count_of_phones, +): + area = CustomBroadcastArea( + name='Example', + polygons=[polygon], + ) + + assert sorted( + overlap.name for overlap in area.overlapping_areas + ) == expected_possible_overlaps + + assert close_enough(area.count_of_phones, expected_count_of_phones) diff --git a/tests/app/main/views/test_broadcast.py b/tests/app/main/views/test_broadcast.py index 4b3f4511e..58671a4ee 100644 --- a/tests/app/main/views/test_broadcast.py +++ b/tests/app/main/views/test_broadcast.py @@ -8,6 +8,7 @@ from flask import url_for from freezegun import freeze_time from tests import broadcast_message_json, sample_uuid, user_json +from tests.app.broadcast_areas.custom_polygons import BRISTOL, SKYE from tests.conftest import SERVICE_ONE_ID, normalize_spaces sample_uuid = sample_uuid() @@ -727,11 +728,42 @@ def test_preview_broadcast_areas_page( ] == estimates +@pytest.mark.parametrize('polygons, expected_list_items', ( + ( + [ + [[1, 2], [3, 4], [5, 6]], + [[7, 8], [9, 10], [11, 12]], + ], + [ + 'An area of 722.3 square miles Will get the alert', + 'An extra area of 1,498.5 square miles is Likely to get the alert', + 'Unknown number of phones', + ] + ), + ( + [BRISTOL], + [ + 'An area of 6.6 square miles Will get the alert', + 'An extra area of 6.4 square miles is Likely to get the alert', + '70,000 to 100,000 phones', + ] + ), + ( + [SKYE], + [ + 'An area of 3,205.0 square miles Will get the alert', + 'An extra area of 763.4 square miles is Likely to get the alert', + '4,000 to 5,000 phones', + ] + ), +)) def test_preview_broadcast_areas_page_with_custom_polygons( mocker, client_request, service_one, fake_uuid, + polygons, + expected_list_items, ): service_one['permissions'] += ['broadcast'] mocker.patch( @@ -743,10 +775,7 @@ def test_preview_broadcast_areas_page_with_custom_polygons( service_id=SERVICE_ONE_ID, status='draft', areas=['Area one', 'Area two', 'Area three'], - simple_polygons=[ - [[1, 2], [3, 4], [5, 6]], - [[7, 8], [9, 10], [11, 12]], - ], + simple_polygons=polygons, ), ) page = client_request.get( @@ -767,11 +796,7 @@ def test_preview_broadcast_areas_page_with_custom_polygons( assert [ normalize_spaces(item.text) for item in page.select('ul li.area-list-key') - ] == [ - 'An area of 722.3 square miles Will get the alert', - 'An extra area of 1,498.5 square miles is Likely to get the alert', - 'Unknown number of phones', - ] + ] == expected_list_items @pytest.mark.parametrize('areas, expected_list', (