Merge pull request #3844 from alphagov/calculate-phones-custom-areas

Calculate the estimated number of phones in an arbitrary polygon

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

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'
 )

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):

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

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

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],
+]

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)

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', (