def _cdist(X, Y, euclidean, periodic_phi, phi_col):
if euclidean:
if periodic_phi:
# delta phis (taking into account periodicity)
# aleady guaranteed for values to be in [0, 2pi]
d_phis = np.pi - np.abs(np.pi - np.abs(X[:,phi_col,None] - Y[:,phi_col]))
# split out common case of having only one other dimension
if X.shape[1] == 2:
non_phi_col = 1 - phi_col
d_ys = X[:,non_phi_col,None] - Y[:,non_phi_col]
out = np.sqrt(d_ys**2 + d_phis**2)
# general case
else:
non_phi_cols = [i for i in range(X.shape[1]) if i != phi_col]
d_others2 = (X[:,None,non_phi_cols] - Y[:,non_phi_cols])**2
out = np.sqrt(d_others2.sum(axis=-1) + d_phis**2)
else:
out = np.empty((len(X), len(Y)), dtype=np.double)
_distance_wrap.cdist_euclidean_double_wrap(X, Y, out)
# spherical measure
else:
# add min/max conditions to ensure valid input
out = np.arccos(np.fmax(np.fmin(np.tensordot(X, Y, axes=(1, 1)), 1.0), -1.0))
return out
def _cdist_euclidean(X, Y, periodic_phi, phi_col):
if periodic_phi:
# delta phis (taking into account periodicity)
# aleady guaranteed for values to be in [0, 2pi]
d_phis = np.pi - np.abs(np.pi - np.abs(X[:,phi_col,np.newaxis] - Y[:,phi_col]))
# split out common case of having only one other dimension
if X.shape[1] == 2:
non_phi_col = 1 - phi_col
d_ys = X[:,non_phi_col,np.newaxis] - Y[:,non_phi_col]
out = np.sqrt(d_ys**2 + d_phis**2)
# general case
else:
non_phi_cols = [i for i in range(X.shape[1]) if i != phi_col]
d_others2 = (X[:,np.newaxis,non_phi_cols] - Y[:,non_phi_cols])**2
out = np.sqrt(d_others2.sum(axis=-1) + d_phis**2)
else:
out = np.empty((len(X), len(Y)), dtype=np.double)
_distance_wrap.cdist_euclidean_double_wrap(X, Y, out)
return out
def cdist_(x: np.ndarray) -> np.ndarray:
"""
Fast C wrapper for spacial distance calculations.
The real function is 'cdist' from 'scipy.spatial.distance'.
This function skips a all the checks
performed by scipy.spatial.distance
Equivalent to (but without all the checks):
>>> from scipy.spatial.distance import cdist
>>> cdist(x, x)
"""
m = x.shape[0]
dm = np.empty((m, m), dtype=np.double)
_distance_wrap.cdist_euclidean_double_wrap(x, x, dm)
# assert np.nan_to_num(dm) == dm, 'np.nan found in dm'
return np.nan_to_num(dm)