from warnings import warn
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from scipy.spatial.distance import pdist, squareform
from scipy.special import rel_entr as relative_entropy
try:
from numba import njit, prange
_HAS_NUMBA = True
except ImportError:
_HAS_NUMBA = False
try:
import psutil
_HAS_PSUTIL = True
except ImportError:
_HAS_PSUTIL = False
from ..network import compute_travel_cost_matrix
def _get_distance_matrix_limit():
"""Determine the memory threshold for switching to the streaming path.
Uses 50% of available system RAM when ``psutil`` is available, falling
back to a default of 8 GB.
"""
if _HAS_PSUTIL:
try:
avail = psutil.virtual_memory().available
return max(avail // 2, 1 * 1024**3) # at least 1 GB
except Exception:
pass
return 8 * 1024**3
# Threshold (in bytes) above which the euclidean path falls back from a dense
# distance matrix to streaming cKDTree queries. Defaults to 25% of available
# RAM (via psutil), or 8 GB if psutil is not installed.
_DISTANCE_MATRIX_LIMIT = _get_distance_matrix_limit()
# ---------------------------------------------------------------------------
# Numba-accelerated core (optional, ~12-15× faster than the pure-numpy loop)
# ---------------------------------------------------------------------------
if _HAS_NUMBA:
@njit(parallel=True, cache=True)
def _kl_all_origins_core(df, dist_matrix, total_pop_by_group, total_pop):
"""Compute KL divergence profiles for all origins in parallel.
Returns (all_kl, all_sorted_indices, all_distances, all_pop_covered)
where each is (n, n) — all_kl[i] is the KL profile for origin i,
already sorted by ascending distance.
"""
n = dist_matrix.shape[0]
k = df.shape[1]
all_kl = np.empty((n, n))
all_sorted = np.empty((n, n), dtype=np.int64)
all_dist = np.empty((n, n))
all_pop = np.empty((n, n))
r = total_pop_by_group / total_pop
for i in prange(n):
sorted_indices = np.argsort(dist_matrix[i])
cumul = np.zeros(k)
for idx in range(n):
row = sorted_indices[idx]
for g in range(k):
cumul[g] += df[row, g]
obs_pop = 0.0
for g in range(k):
obs_pop += cumul[g]
kl_sum = 0.0
for g in range(k):
q = cumul[g] / obs_pop
if q > 0:
kl_sum += q * np.log(q / r[g])
all_kl[i, idx] = kl_sum
all_sorted[i, idx] = sorted_indices[idx]
all_dist[i, idx] = dist_matrix[i, sorted_indices[idx]]
all_pop[i, idx] = obs_pop
return all_kl, all_sorted, all_dist, all_pop
@njit(cache=True)
def _kl_single_origin_core(df_row_sorted, r):
"""Compute KL divergence + population_covered for one origin.
Parameters
----------
df_row_sorted : np.ndarray
(n, k) population data already sorted by ascending distance.
r : np.ndarray
(k,) global proportions.
Returns
-------
kl : np.ndarray (n,)
KL divergence at each cumulative radius.
pop_covered : np.ndarray (n,)
Cumulative population covered at each radius.
"""
n, k = df_row_sorted.shape
kl = np.empty(n)
pop_covered = np.empty(n)
cumul = np.zeros(k)
for idx in range(n):
for g in range(k):
cumul[g] += df_row_sorted[idx, g]
obs_pop = 0.0
for g in range(k):
obs_pop += cumul[g]
kl_sum = 0.0
for g in range(k):
q = cumul[g] / obs_pop
if q > 0:
kl_sum += q * np.log(q / r[g])
kl[idx] = kl_sum
pop_covered[idx] = obs_pop
return kl, pop_covered
def _kl_profile_for_origin(sorted_indices, distances, df, indices, i):
"""Compute the KL divergence profile for a single origin.
Uses the numba-accelerated single-origin kernel when available.
Parameters
----------
sorted_indices : np.ndarray
Unit indices sorted by ascending distance from the origin.
distances : np.ndarray
Corresponding distances (already sorted ascending).
df : np.ndarray
(n, k) array of group population counts.
indices : pd.Index
Original observation index for labelling.
i : int
Positional index of the origin.
Returns
-------
pd.DataFrame
Divergence profile for this origin.
"""
df_sorted = df[sorted_indices].astype(np.float64)
total_pop = float(df.sum())
r = df.sum(axis=0).astype(np.float64) / total_pop
if _HAS_NUMBA:
kl_divergence, pop_covered = _kl_single_origin_core(df_sorted, r)
else:
cumul_pop_by_group = np.cumsum(df_sorted, axis=0)
obs_cumul_pop = np.sum(cumul_pop_by_group, axis=1)[:, np.newaxis]
q_cumul_proportions = cumul_pop_by_group / obs_cumul_pop
r_total_proportions = r[np.newaxis, :]
kl_divergence = relative_entropy(q_cumul_proportions, r_total_proportions).sum(
axis=1
)
pop_covered = obs_cumul_pop.sum(axis=1)
return pd.DataFrame().from_dict(
dict(
observation=indices[i],
distance=distances,
divergence=kl_divergence,
population_covered=pop_covered,
)
)
def _compute_all_profiles(dist_matrix, df, indices):
"""Compute KL divergence profiles for all origins at once.
Uses the numba-accelerated parallel kernel when available (~12-15× faster),
falling back to the pure-numpy per-origin loop otherwise.
Parameters
----------
dist_matrix : np.ndarray
(n, n) dense distance matrix.
df : np.ndarray
(n, k) array of group population counts.
indices : pd.Index
Original observation index for labelling.
Returns
-------
list[pd.DataFrame]
One divergence-profile DataFrame per origin.
"""
n = len(df)
total_pop_by_group = df.sum(axis=0).astype(np.float64)
total_pop = float(df.sum())
if _HAS_NUMBA:
all_kl, all_sorted, all_dist, all_pop = _kl_all_origins_core(
df.astype(np.float64), dist_matrix, total_pop_by_group, total_pop
)
results = []
for i in range(n):
results.append(
pd.DataFrame().from_dict(
dict(
observation=indices[i],
distance=all_dist[i],
divergence=all_kl[i],
population_covered=all_pop[i],
)
)
)
return results
else:
results = []
for i in range(n):
sorted_indices = np.argsort(dist_matrix[i])
results.append(
_kl_profile_for_origin(
sorted_indices,
dist_matrix[i][sorted_indices],
df,
indices,
i,
)
)
return results
[docs]
def compute_divergence_profiles(
gdf, groups, metric="euclidean", network=None, distance_matrix=None
):
"""
A segregation metric using Kullback-Leiber (KL) divergence to quantify the
difference in the population characteristics between (1) an area and (2) the total population.
Parameters
----------
data : pandas.DataFrame or geopandas.GeoDataFrame, required
dataframe or geodataframe if spatial index holding data for location of interest
groups : list, required
list of columns on dataframe holding population totals for each group
metric : str (optional; 'euclidean' by default)
Distance metric for calculating pairwise distances,
Accepts any inputs to `scipy.spatial.distance.pdist`.
Ignored if passing a network or distance matrix
network: pandarm.Network object (optional, None by default)
A pandarm Network object used to compute distance between observations
distance_matrix: numpy.array (optional; None by default)
numpy array of distances between observations in the dataset
Returns
----------
aux : geopandas.GeoDataFrame
geodataframe of the KL divergence measure, between the aggregated population and the
total population, will converge to zero for the final row of each
observation to represent that the total population is covered.
population_covered : the population count within the aggregated population.
Returns a concatenated object of Pandas dataframes. Each dataframe contains a
set of divergence levels between an area and the total population. These areas
become consecutively larger, starting from a single location and aggregating
outward from this location, until the area represents the total population.
Thus, together the divergence levels within a dataframe represent a profile
of divergence from an area. The concatenated object is the collection of these
divergence profiles for every areas within the total population.
Notes
-----
When ``numba`` is installed, the per-origin KL divergence computation is
JIT-compiled and parallelized across CPU cores, yielding ~12-15× speedup
over the pure-numpy loop. If ``numba`` is not available, the function
falls back to the numpy implementation automatically.
For the default euclidean metric, the implementation uses a dense distance
matrix via ``scipy.spatial.distance.pdist`` (fast, vectorized C) when the
matrix fits comfortably in memory. For very large datasets where the dense
matrix would be impractical, it falls back to a streaming ``cKDTree`` path
that keeps peak memory at *O(n)* per iteration. The threshold defaults to
50% of available system RAM (via ``psutil``), or 8 GB if ``psutil`` is not
installed. It can be overridden by setting
``segregation.dynamics.divergence_profile._DISTANCE_MATRIX_LIMIT``.
"""
# Store the observation index to return with the results
indices = gdf.index.copy()
centroids = gdf.geometry.centroid
df = gdf[groups].values
n = len(df)
coordinates = np.column_stack((centroids.x, centroids.y))
# Preparing list for results
results = []
if network:
# --- Network path: dense matrix from pandarm ---
if metric != "network":
warn(
f"metric set to {metric} but a pandarm.Network object was passed. Using network distances instead"
"If you wish to use a scipy distance matrix, do not include a `network` argument`",
stacklevel=2,
)
dist_matrix = compute_travel_cost_matrix(gdf, gdf, network).values
results = _compute_all_profiles(dist_matrix, df, indices)
elif distance_matrix is not None:
# --- Precomputed matrix path ---
if hasattr(distance_matrix, "tocsr"):
# A divergence profile aggregates outward from each origin until it
# covers the *total* population, so it needs a distance from every
# origin to every unit. A sparse matrix leaves some pairs
# undefined, which truncates the profile at a different point for
# each origin and makes the resulting coefficients incomparable.
raise TypeError(
"sparse distance matrices are not supported: a divergence "
"profile requires a distance from every origin to every "
"observation. Pass a dense numpy array as `distance_matrix`"
)
if metric != "precomputed":
warn(
f"metric set to {metric} but a distance_matrix argument was passed. Using precomputed distances instead",
stacklevel=2,
)
results = _compute_all_profiles(distance_matrix, df, indices)
elif metric == "euclidean":
# --- Euclidean path ---
# Use the fast dense-matrix path (pdist) when the matrix fits in
# memory; fall back to streaming cKDTree for very large datasets.
matrix_bytes = n * n * 8
if matrix_bytes <= _DISTANCE_MATRIX_LIMIT:
dist_matrix = squareform(pdist(coordinates, metric=metric))
results = _compute_all_profiles(dist_matrix, df, indices)
else:
# Streaming path: query in chunks to amortize Python→C overhead
# while keeping peak memory at O(chunk × n) not O(n²).
tree = cKDTree(coordinates)
chunk_size = min(n, 5000)
for start in range(0, n, chunk_size):
end = min(start + chunk_size, n)
_distances, _sorted_indices = tree.query(coordinates[start:end], k=n)
for j in range(end - start):
i = start + j
results.append(
_kl_profile_for_origin(
_sorted_indices[j], _distances[j], df, indices, i
)
)
else:
# --- Non-euclidean scipy metric: fall back to dense matrix ---
dist_matrix = squareform(pdist(coordinates, metric=metric))
results = _compute_all_profiles(dist_matrix, df, indices)
aux = pd.concat(results)
return aux