distance.pyi 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. from typing import (overload, Any, SupportsFloat, Literal, Protocol, SupportsIndex)
  2. import numpy as np
  3. from numpy.typing import ArrayLike, NDArray
  4. # Anything that can be parsed by `np.float64.__init__` and is thus
  5. # compatible with `ndarray.__setitem__` (for a float64 array)
  6. _FloatValue = None | str | bytes | SupportsFloat | SupportsIndex
  7. class _MetricCallback1(Protocol):
  8. def __call__(
  9. self, __XA: NDArray[Any], __XB: NDArray[Any]
  10. ) -> _FloatValue: ...
  11. class _MetricCallback2(Protocol):
  12. def __call__(
  13. self, __XA: NDArray[Any], __XB: NDArray[Any], **kwargs: Any
  14. ) -> _FloatValue: ...
  15. # TODO: Use a single protocol with a parameter specification variable
  16. # once available (PEP 612)
  17. _MetricCallback = _MetricCallback1 | _MetricCallback2
  18. _MetricKind = Literal[
  19. 'braycurtis',
  20. 'canberra',
  21. 'chebychev', 'chebyshev', 'cheby', 'cheb', 'ch',
  22. 'cityblock', 'cblock', 'cb', 'c',
  23. 'correlation', 'co',
  24. 'cosine', 'cos',
  25. 'dice',
  26. 'euclidean', 'euclid', 'eu', 'e',
  27. 'hamming', 'hamm', 'ha', 'h',
  28. 'minkowski', 'mi', 'm', 'pnorm',
  29. 'jaccard', 'jacc', 'ja', 'j',
  30. 'jensenshannon', 'js',
  31. 'mahalanobis', 'mahal', 'mah',
  32. 'rogerstanimoto',
  33. 'russellrao',
  34. 'seuclidean', 'se', 's',
  35. 'sokalsneath',
  36. 'sqeuclidean', 'sqe', 'sqeuclid',
  37. 'yule',
  38. ]
  39. # Function annotations
  40. def braycurtis(
  41. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  42. ) -> np.float64: ...
  43. def canberra(
  44. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  45. ) -> np.float64: ...
  46. # TODO: Add `metric`-specific overloads
  47. # Returns a float64 or float128 array, depending on the input dtype
  48. @overload
  49. def cdist(
  50. XA: ArrayLike,
  51. XB: ArrayLike,
  52. metric: _MetricKind = ...,
  53. *,
  54. out: None | NDArray[np.floating[Any]] = ...,
  55. p: float = ...,
  56. w: ArrayLike | None = ...,
  57. V: ArrayLike | None = ...,
  58. VI: ArrayLike | None = ...,
  59. ) -> NDArray[np.floating[Any]]: ...
  60. @overload
  61. def cdist(
  62. XA: ArrayLike,
  63. XB: ArrayLike,
  64. metric: _MetricCallback,
  65. *,
  66. out: None | NDArray[np.floating[Any]] = ...,
  67. **kwargs: Any,
  68. ) -> NDArray[np.floating[Any]]: ...
  69. # TODO: Wait for dtype support; the return type is
  70. # dependent on the input arrays dtype
  71. def chebyshev(
  72. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  73. ) -> Any: ...
  74. # TODO: Wait for dtype support; the return type is
  75. # dependent on the input arrays dtype
  76. def cityblock(
  77. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  78. ) -> Any: ...
  79. def correlation(
  80. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ..., centered: bool = ...
  81. ) -> np.float64: ...
  82. def cosine(
  83. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  84. ) -> np.float64: ...
  85. def dice(
  86. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  87. ) -> float: ...
  88. def directed_hausdorff(
  89. u: ArrayLike, v: ArrayLike, seed: int | None = ...
  90. ) -> tuple[float, int, int]: ...
  91. def euclidean(
  92. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  93. ) -> float: ...
  94. def hamming(
  95. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  96. ) -> np.float64: ...
  97. def is_valid_dm(
  98. D: ArrayLike,
  99. tol: float = ...,
  100. throw: bool = ...,
  101. name: str | None = ...,
  102. warning: bool = ...,
  103. ) -> bool: ...
  104. def is_valid_y(
  105. y: ArrayLike,
  106. warning: bool = ...,
  107. throw: bool = ...,
  108. name: str | None = ...,
  109. ) -> bool: ...
  110. def jaccard(
  111. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  112. ) -> np.float64: ...
  113. def jensenshannon(
  114. p: ArrayLike, q: ArrayLike, base: float | None = ...
  115. ) -> np.float64: ...
  116. def mahalanobis(
  117. u: ArrayLike, v: ArrayLike, VI: ArrayLike
  118. ) -> np.float64: ...
  119. def minkowski(
  120. u: ArrayLike, v: ArrayLike, p: float = ..., w: ArrayLike | None = ...
  121. ) -> float: ...
  122. def num_obs_dm(d: ArrayLike) -> int: ...
  123. def num_obs_y(Y: ArrayLike) -> int: ...
  124. # TODO: Add `metric`-specific overloads
  125. @overload
  126. def pdist(
  127. X: ArrayLike,
  128. metric: _MetricKind = ...,
  129. *,
  130. out: None | NDArray[np.floating[Any]] = ...,
  131. p: float = ...,
  132. w: ArrayLike | None = ...,
  133. V: ArrayLike | None = ...,
  134. VI: ArrayLike | None = ...,
  135. ) -> NDArray[np.floating[Any]]: ...
  136. @overload
  137. def pdist(
  138. X: ArrayLike,
  139. metric: _MetricCallback,
  140. *,
  141. out: None | NDArray[np.floating[Any]] = ...,
  142. **kwargs: Any,
  143. ) -> NDArray[np.floating[Any]]: ...
  144. def seuclidean(
  145. u: ArrayLike, v: ArrayLike, V: ArrayLike
  146. ) -> float: ...
  147. def sokalsneath(
  148. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  149. ) -> np.float64: ...
  150. def sqeuclidean(
  151. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  152. ) -> np.float64: ...
  153. def squareform(
  154. X: ArrayLike,
  155. force: Literal["no", "tomatrix", "tovector"] = ...,
  156. checks: bool = ...,
  157. ) -> NDArray[Any]: ...
  158. def rogerstanimoto(
  159. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  160. ) -> float: ...
  161. def russellrao(
  162. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  163. ) -> float: ...
  164. def yule(
  165. u: ArrayLike, v: ArrayLike, w: ArrayLike | None = ...
  166. ) -> float: ...