wehub-resource-sync
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
from hypothesis import strategies as st
|
|
import numpy as np
|
|
from scipy.sparse import csr_matrix
|
|
|
|
|
|
@st.composite
|
|
def knn_graph_strategy(draw, num_samples, k_neighbors, min_distance=0.0, max_distance=100.0):
|
|
"""
|
|
Generate a K-nearest neighbors (KNN) graph based on the given parameters.
|
|
|
|
Parameters
|
|
----------
|
|
draw: A function used to draw values from search strategies.
|
|
|
|
num_samples (int or SearchStrategy): The number of samples in the graph.
|
|
If a SearchStrategy is provided, a value will be drawn from it.
|
|
|
|
k_neighbors (int or SearchStrategy): The number of nearest neighbors to consider for each sample.
|
|
If a SearchStrategy is provided, a value will be drawn from it.
|
|
|
|
Returns
|
|
-------
|
|
knn_graph : csr_matrix
|
|
The KNN graph represented as a sparse matrix.
|
|
|
|
Notes
|
|
-----
|
|
- The KNN graph is generated based on a symmetric distance matrix.
|
|
- The distance matrix is computed using randomly generated upper triangle values.
|
|
- The diagonal of the distance matrix is set to infinity to avoid selecting a point as its own neighbor.
|
|
- The K-nearest neighbors are computed based on the distance matrix.
|
|
- The resulting KNN graph is returned as a sparse matrix in csr format.
|
|
- The number of samples must be greater than the number of neighbors.
|
|
- The KNN graph is not guaranteed to be connected (i.e. there may be isolated subgraphs).
|
|
- The KNN graph is a directed graph (i.e. the edges are not symmetric).
|
|
- The neighbors are sorted by distance in the CSR-formatted sparse matrix,
|
|
so the first neighbor is the closest neighbor.
|
|
"""
|
|
# If the argument is a strategy, draw a value from it.
|
|
if isinstance(num_samples, st.SearchStrategy):
|
|
num_samples = draw(num_samples)
|
|
|
|
if isinstance(k_neighbors, st.SearchStrategy):
|
|
k_neighbors = draw(k_neighbors)
|
|
|
|
# Generate a symmetric distance matrix
|
|
upper_triangle = [
|
|
draw(
|
|
st.lists(
|
|
st.floats(
|
|
min_value=min_distance,
|
|
max_value=max_distance,
|
|
allow_nan=False,
|
|
allow_infinity=False,
|
|
allow_subnormal=False,
|
|
),
|
|
min_size=i,
|
|
max_size=i,
|
|
unique=True,
|
|
)
|
|
)
|
|
for i in range(1, num_samples + 1)
|
|
]
|
|
|
|
distance_matrix = np.zeros((num_samples, num_samples))
|
|
for i, row in enumerate(upper_triangle):
|
|
distance_matrix[i, : i + 1] = row
|
|
distance_matrix[: i + 1, i] = row
|
|
|
|
np.fill_diagonal(
|
|
distance_matrix, np.inf
|
|
) # To ensure we don't select a point as its own neighbor
|
|
|
|
# Compute k-nearest neighbors based on the distance matrix
|
|
sorted_indices = np.argsort(distance_matrix, axis=1)
|
|
kneighbor_indices = sorted_indices[:, :k_neighbors]
|
|
kneighbor_distances = np.array(
|
|
[distance_matrix[i, kneighbor_indices[i]] for i in range(num_samples)]
|
|
)
|
|
|
|
knn_graph = csr_matrix(
|
|
(
|
|
kneighbor_distances.flatten(),
|
|
kneighbor_indices.flatten(),
|
|
np.arange(0, (kneighbor_distances.shape[0] * k_neighbors + 1), k_neighbors),
|
|
),
|
|
shape=(num_samples, num_samples),
|
|
)
|
|
return knn_graph
|