array.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261
  1. from __future__ import annotations
  2. from itertools import chain
  3. from typing import TYPE_CHECKING
  4. import numpy as np
  5. from contourpy.typecheck import check_code_array, check_offset_array, check_point_array
  6. from contourpy.types import CLOSEPOLY, LINETO, MOVETO, code_dtype, offset_dtype, point_dtype
  7. if TYPE_CHECKING:
  8. import contourpy._contourpy as cpy
  9. def codes_from_offsets(offsets: cpy.OffsetArray) -> cpy.CodeArray:
  10. """Determine codes from offsets, assuming they all correspond to closed polygons.
  11. """
  12. check_offset_array(offsets)
  13. n = offsets[-1]
  14. codes = np.full(n, LINETO, dtype=code_dtype)
  15. codes[offsets[:-1]] = MOVETO
  16. codes[offsets[1:] - 1] = CLOSEPOLY
  17. return codes
  18. def codes_from_offsets_and_points(
  19. offsets: cpy.OffsetArray,
  20. points: cpy.PointArray,
  21. ) -> cpy.CodeArray:
  22. """Determine codes from offsets and points, using the equality of the start and end points of
  23. each line to determine if lines are closed or not.
  24. """
  25. check_offset_array(offsets)
  26. check_point_array(points)
  27. codes = np.full(len(points), LINETO, dtype=code_dtype)
  28. codes[offsets[:-1]] = MOVETO
  29. end_offsets = offsets[1:] - 1
  30. closed = np.all(points[offsets[:-1]] == points[end_offsets], axis=1)
  31. codes[end_offsets[closed]] = CLOSEPOLY
  32. return codes
  33. def codes_from_points(points: cpy.PointArray) -> cpy.CodeArray:
  34. """Determine codes for a single line, using the equality of the start and end points to
  35. determine if the line is closed or not.
  36. """
  37. check_point_array(points)
  38. n = len(points)
  39. codes = np.full(n, LINETO, dtype=code_dtype)
  40. codes[0] = MOVETO
  41. if np.all(points[0] == points[-1]):
  42. codes[-1] = CLOSEPOLY
  43. return codes
  44. def concat_codes(list_of_codes: list[cpy.CodeArray]) -> cpy.CodeArray:
  45. """Concatenate a list of codes arrays into a single code array.
  46. """
  47. if not list_of_codes:
  48. raise ValueError("Empty list passed to concat_codes")
  49. return np.concatenate(list_of_codes, dtype=code_dtype)
  50. def concat_codes_or_none(list_of_codes_or_none: list[cpy.CodeArray | None]) -> cpy.CodeArray | None:
  51. """Concatenate a list of codes arrays or None into a single code array or None.
  52. """
  53. list_of_codes = [codes for codes in list_of_codes_or_none if codes is not None]
  54. if list_of_codes:
  55. return concat_codes(list_of_codes)
  56. else:
  57. return None
  58. def concat_offsets(list_of_offsets: list[cpy.OffsetArray]) -> cpy.OffsetArray:
  59. """Concatenate a list of offsets arrays into a single offset array.
  60. """
  61. if not list_of_offsets:
  62. raise ValueError("Empty list passed to concat_offsets")
  63. n = len(list_of_offsets)
  64. cumulative = np.cumsum([offsets[-1] for offsets in list_of_offsets], dtype=offset_dtype)
  65. ret: cpy.OffsetArray = np.concatenate(
  66. (list_of_offsets[0], *(list_of_offsets[i+1][1:] + cumulative[i] for i in range(n-1))),
  67. dtype=offset_dtype,
  68. )
  69. return ret
  70. def concat_offsets_or_none(
  71. list_of_offsets_or_none: list[cpy.OffsetArray | None],
  72. ) -> cpy.OffsetArray | None:
  73. """Concatenate a list of offsets arrays or None into a single offset array or None.
  74. """
  75. list_of_offsets = [offsets for offsets in list_of_offsets_or_none if offsets is not None]
  76. if list_of_offsets:
  77. return concat_offsets(list_of_offsets)
  78. else:
  79. return None
  80. def concat_points(list_of_points: list[cpy.PointArray]) -> cpy.PointArray:
  81. """Concatenate a list of point arrays into a single point array.
  82. """
  83. if not list_of_points:
  84. raise ValueError("Empty list passed to concat_points")
  85. return np.concatenate(list_of_points, dtype=point_dtype)
  86. def concat_points_or_none(
  87. list_of_points_or_none: list[cpy.PointArray | None],
  88. ) -> cpy.PointArray | None:
  89. """Concatenate a list of point arrays or None into a single point array or None.
  90. """
  91. list_of_points = [points for points in list_of_points_or_none if points is not None]
  92. if list_of_points:
  93. return concat_points(list_of_points)
  94. else:
  95. return None
  96. def concat_points_or_none_with_nan(
  97. list_of_points_or_none: list[cpy.PointArray | None],
  98. ) -> cpy.PointArray | None:
  99. """Concatenate a list of points or None into a single point array or None, with NaNs used to
  100. separate each line.
  101. """
  102. list_of_points = [points for points in list_of_points_or_none if points is not None]
  103. if list_of_points:
  104. return concat_points_with_nan(list_of_points)
  105. else:
  106. return None
  107. def concat_points_with_nan(list_of_points: list[cpy.PointArray]) -> cpy.PointArray:
  108. """Concatenate a list of points into a single point array with NaNs used to separate each line.
  109. """
  110. if not list_of_points:
  111. raise ValueError("Empty list passed to concat_points_with_nan")
  112. if len(list_of_points) == 1:
  113. return list_of_points[0]
  114. else:
  115. nan_spacer = np.full((1, 2), np.nan, dtype=point_dtype)
  116. list_of_points = [list_of_points[0],
  117. *list(chain(*((nan_spacer, x) for x in list_of_points[1:])))]
  118. return concat_points(list_of_points)
  119. def insert_nan_at_offsets(points: cpy.PointArray, offsets: cpy.OffsetArray) -> cpy.PointArray:
  120. """Insert NaNs into a point array at locations specified by an offset array.
  121. """
  122. check_point_array(points)
  123. check_offset_array(offsets)
  124. if len(offsets) <= 2:
  125. return points
  126. else:
  127. nan_spacer = np.array([np.nan, np.nan], dtype=point_dtype)
  128. # Convert offsets to int64 to avoid numpy error when mixing signed and unsigned ints.
  129. return np.insert(points, offsets[1:-1].astype(np.int64), nan_spacer, axis=0)
  130. def offsets_from_codes(codes: cpy.CodeArray) -> cpy.OffsetArray:
  131. """Determine offsets from codes using locations of MOVETO codes.
  132. """
  133. check_code_array(codes)
  134. return np.append(np.nonzero(codes == MOVETO)[0], len(codes)).astype(offset_dtype)
  135. def offsets_from_lengths(list_of_points: list[cpy.PointArray]) -> cpy.OffsetArray:
  136. """Determine offsets from lengths of point arrays.
  137. """
  138. if not list_of_points:
  139. raise ValueError("Empty list passed to offsets_from_lengths")
  140. return np.cumsum([0] + [len(line) for line in list_of_points], dtype=offset_dtype)
  141. def outer_offsets_from_list_of_codes(list_of_codes: list[cpy.CodeArray]) -> cpy.OffsetArray:
  142. """Determine outer offsets from codes using locations of MOVETO codes.
  143. """
  144. if not list_of_codes:
  145. raise ValueError("Empty list passed to outer_offsets_from_list_of_codes")
  146. return np.cumsum([0] + [np.count_nonzero(codes == MOVETO) for codes in list_of_codes],
  147. dtype=offset_dtype)
  148. def outer_offsets_from_list_of_offsets(list_of_offsets: list[cpy.OffsetArray]) -> cpy.OffsetArray:
  149. """Determine outer offsets from a list of offsets.
  150. """
  151. if not list_of_offsets:
  152. raise ValueError("Empty list passed to outer_offsets_from_list_of_offsets")
  153. return np.cumsum([0] + [len(offsets)-1 for offsets in list_of_offsets], dtype=offset_dtype)
  154. def remove_nan(points: cpy.PointArray) -> tuple[cpy.PointArray, cpy.OffsetArray]:
  155. """Remove NaN from a points array, also return the offsets corresponding to the NaN removed.
  156. """
  157. check_point_array(points)
  158. nan_offsets = np.nonzero(np.isnan(points[:, 0]))[0]
  159. if len(nan_offsets) == 0:
  160. return points, np.array([0, len(points)], dtype=offset_dtype)
  161. else:
  162. points = np.delete(points, nan_offsets, axis=0)
  163. nan_offsets -= np.arange(len(nan_offsets))
  164. offsets: cpy.OffsetArray = np.empty(len(nan_offsets)+2, dtype=offset_dtype)
  165. offsets[0] = 0
  166. offsets[1:-1] = nan_offsets
  167. offsets[-1] = len(points)
  168. return points, offsets
  169. def split_codes_by_offsets(codes: cpy.CodeArray, offsets: cpy.OffsetArray) -> list[cpy.CodeArray]:
  170. """Split a code array at locations specified by an offset array into a list of code arrays.
  171. """
  172. check_code_array(codes)
  173. check_offset_array(offsets)
  174. if len(offsets) > 2:
  175. return np.split(codes, offsets[1:-1])
  176. else:
  177. return [codes]
  178. def split_points_by_offsets(
  179. points: cpy.PointArray,
  180. offsets: cpy.OffsetArray,
  181. ) -> list[cpy.PointArray]:
  182. """Split a point array at locations specified by an offset array into a list of point arrays.
  183. """
  184. check_point_array(points)
  185. check_offset_array(offsets)
  186. if len(offsets) > 2:
  187. return np.split(points, offsets[1:-1])
  188. else:
  189. return [points]
  190. def split_points_at_nan(points: cpy.PointArray) -> list[cpy.PointArray]:
  191. """Split a points array at NaNs into a list of point arrays.
  192. """
  193. check_point_array(points)
  194. nan_offsets = np.nonzero(np.isnan(points[:, 0]))[0]
  195. if len(nan_offsets) == 0:
  196. return [points]
  197. else:
  198. nan_offsets = np.concatenate(([-1], nan_offsets, [len(points)])) # type: ignore[arg-type]
  199. return [points[s+1:e] for s, e in zip(nan_offsets[:-1], nan_offsets[1:])]