-
Notifications
You must be signed in to change notification settings - Fork 2
/
geometry.py
35 lines (31 loc) · 1 KB
/
geometry.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from collections import defaultdict
import numpy as np
import tqdm.auto as tqdm
from permacache import permacache, stable_hash
from urbanstats.geometry.categorize_coordinates import categorize
from urbanstats.geometry.ellipse import Ellipse
@permacache(
"population_density/geometry/locate_blocks_2",
key_function=dict(
coordinates=stable_hash,
population=stable_hash,
),
multiprocess_safe=True,
)
def locate_blocks(*, coordinates, population, radius=1):
categories = categorize(coordinates)
result = defaultdict(list)
for i in tqdm.trange(coordinates.shape[0]):
result[tuple(categories[i])].append(i)
result = {
cat: dict(indices=np.array(res), coordinates=coordinates[np.array(res)])
for cat, res in tqdm.tqdm(result.items())
}
return np.array(
[
population[Ellipse(radius, *coord).find_neighbors(result, coordinates)].sum(
0
)
for coord in tqdm.tqdm(coordinates)
]
)