dechunk.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207
  1. from __future__ import annotations
  2. from typing import TYPE_CHECKING, cast
  3. from contourpy._contourpy import FillType, LineType
  4. from contourpy.array import (
  5. concat_codes_or_none,
  6. concat_offsets_or_none,
  7. concat_points_or_none,
  8. concat_points_or_none_with_nan,
  9. )
  10. from contourpy.enum_util import as_fill_type, as_line_type
  11. from contourpy.typecheck import check_filled, check_lines
  12. if TYPE_CHECKING:
  13. import contourpy._contourpy as cpy
  14. def dechunk_filled(filled: cpy.FillReturn, fill_type: FillType | str) -> cpy.FillReturn:
  15. """Return the specified filled contours with chunked data moved into the first chunk.
  16. Filled contours that are not chunked (``FillType.OuterCode`` and ``FillType.OuterOffset``) and
  17. those that are but only contain a single chunk are returned unmodified. Individual polygons are
  18. unchanged, they are not geometrically combined.
  19. Args:
  20. filled (sequence of arrays): Filled contour data, such as returned by
  21. :meth:`.ContourGenerator.filled`.
  22. fill_type (FillType or str): Type of :meth:`~.ContourGenerator.filled` as enum or string
  23. equivalent.
  24. Return:
  25. Filled contours in a single chunk.
  26. .. versionadded:: 1.2.0
  27. """
  28. fill_type = as_fill_type(fill_type)
  29. if fill_type in (FillType.OuterCode, FillType.OuterOffset):
  30. # No-op if fill_type is not chunked.
  31. return filled
  32. check_filled(filled, fill_type)
  33. if len(filled[0]) < 2:
  34. # No-op if just one chunk.
  35. return filled
  36. if TYPE_CHECKING:
  37. filled = cast(cpy.FillReturn_Chunk, filled)
  38. points = concat_points_or_none(filled[0])
  39. if fill_type == FillType.ChunkCombinedCode:
  40. if TYPE_CHECKING:
  41. filled = cast(cpy.FillReturn_ChunkCombinedCode, filled)
  42. if points is None:
  43. ret1: cpy.FillReturn_ChunkCombinedCode = ([None], [None])
  44. else:
  45. ret1 = ([points], [concat_codes_or_none(filled[1])])
  46. return ret1
  47. elif fill_type == FillType.ChunkCombinedOffset:
  48. if TYPE_CHECKING:
  49. filled = cast(cpy.FillReturn_ChunkCombinedOffset, filled)
  50. if points is None:
  51. ret2: cpy.FillReturn_ChunkCombinedOffset = ([None], [None])
  52. else:
  53. ret2 = ([points], [concat_offsets_or_none(filled[1])])
  54. return ret2
  55. elif fill_type == FillType.ChunkCombinedCodeOffset:
  56. if TYPE_CHECKING:
  57. filled = cast(cpy.FillReturn_ChunkCombinedCodeOffset, filled)
  58. if points is None:
  59. ret3: cpy.FillReturn_ChunkCombinedCodeOffset = ([None], [None], [None])
  60. else:
  61. outer_offsets = concat_offsets_or_none(filled[2])
  62. ret3 = ([points], [concat_codes_or_none(filled[1])], [outer_offsets])
  63. return ret3
  64. elif fill_type == FillType.ChunkCombinedOffsetOffset:
  65. if TYPE_CHECKING:
  66. filled = cast(cpy.FillReturn_ChunkCombinedOffsetOffset, filled)
  67. if points is None:
  68. ret4: cpy.FillReturn_ChunkCombinedOffsetOffset = ([None], [None], [None])
  69. else:
  70. outer_offsets = concat_offsets_or_none(filled[2])
  71. ret4 = ([points], [concat_offsets_or_none(filled[1])], [outer_offsets])
  72. return ret4
  73. else:
  74. raise ValueError(f"Invalid FillType {fill_type}")
  75. def dechunk_lines(lines: cpy.LineReturn, line_type: LineType | str) -> cpy.LineReturn:
  76. """Return the specified contour lines with chunked data moved into the first chunk.
  77. Contour lines that are not chunked (``LineType.Separate`` and ``LineType.SeparateCode``) and
  78. those that are but only contain a single chunk are returned unmodified. Individual lines are
  79. unchanged, they are not geometrically combined.
  80. Args:
  81. lines (sequence of arrays): Contour line data, such as returned by
  82. :meth:`.ContourGenerator.lines`.
  83. line_type (LineType or str): Type of :meth:`~.ContourGenerator.lines` as enum or string
  84. equivalent.
  85. Return:
  86. Contour lines in a single chunk.
  87. .. versionadded:: 1.2.0
  88. """
  89. line_type = as_line_type(line_type)
  90. if line_type in (LineType.Separate, LineType.SeparateCode):
  91. # No-op if line_type is not chunked.
  92. return lines
  93. check_lines(lines, line_type)
  94. if len(lines[0]) < 2:
  95. # No-op if just one chunk.
  96. return lines
  97. if TYPE_CHECKING:
  98. lines = cast(cpy.LineReturn_Chunk, lines)
  99. if line_type == LineType.ChunkCombinedCode:
  100. if TYPE_CHECKING:
  101. lines = cast(cpy.LineReturn_ChunkCombinedCode, lines)
  102. points = concat_points_or_none(lines[0])
  103. if points is None:
  104. ret1: cpy.LineReturn_ChunkCombinedCode = ([None], [None])
  105. else:
  106. ret1 = ([points], [concat_codes_or_none(lines[1])])
  107. return ret1
  108. elif line_type == LineType.ChunkCombinedOffset:
  109. if TYPE_CHECKING:
  110. lines = cast(cpy.LineReturn_ChunkCombinedOffset, lines)
  111. points = concat_points_or_none(lines[0])
  112. if points is None:
  113. ret2: cpy.LineReturn_ChunkCombinedOffset = ([None], [None])
  114. else:
  115. ret2 = ([points], [concat_offsets_or_none(lines[1])])
  116. return ret2
  117. elif line_type == LineType.ChunkCombinedNan:
  118. if TYPE_CHECKING:
  119. lines = cast(cpy.LineReturn_ChunkCombinedNan, lines)
  120. points = concat_points_or_none_with_nan(lines[0])
  121. ret3: cpy.LineReturn_ChunkCombinedNan = ([points],)
  122. return ret3
  123. else:
  124. raise ValueError(f"Invalid LineType {line_type}")
  125. def dechunk_multi_filled(
  126. multi_filled: list[cpy.FillReturn],
  127. fill_type: FillType | str,
  128. ) -> list[cpy.FillReturn]:
  129. """Return multiple sets of filled contours with chunked data moved into the first chunks.
  130. Filled contours that are not chunked (``FillType.OuterCode`` and ``FillType.OuterOffset``) and
  131. those that are but only contain a single chunk are returned unmodified. Individual polygons are
  132. unchanged, they are not geometrically combined.
  133. Args:
  134. multi_filled (nested sequence of arrays): Filled contour data, such as returned by
  135. :meth:`.ContourGenerator.multi_filled`.
  136. fill_type (FillType or str): Type of :meth:`~.ContourGenerator.filled` as enum or string
  137. equivalent.
  138. Return:
  139. Multiple sets of filled contours in a single chunk.
  140. .. versionadded:: 1.3.0
  141. """
  142. fill_type = as_fill_type(fill_type)
  143. if fill_type in (FillType.OuterCode, FillType.OuterOffset):
  144. # No-op if fill_type is not chunked.
  145. return multi_filled
  146. return [dechunk_filled(filled, fill_type) for filled in multi_filled]
  147. def dechunk_multi_lines(
  148. multi_lines: list[cpy.LineReturn],
  149. line_type: LineType | str,
  150. ) -> list[cpy.LineReturn]:
  151. """Return multiple sets of contour lines with all chunked data moved into the first chunks.
  152. Contour lines that are not chunked (``LineType.Separate`` and ``LineType.SeparateCode``) and
  153. those that are but only contain a single chunk are returned unmodified. Individual lines are
  154. unchanged, they are not geometrically combined.
  155. Args:
  156. multi_lines (nested sequence of arrays): Contour line data, such as returned by
  157. :meth:`.ContourGenerator.multi_lines`.
  158. line_type (LineType or str): Type of :meth:`~.ContourGenerator.lines` as enum or string
  159. equivalent.
  160. Return:
  161. Multiple sets of contour lines in a single chunk.
  162. .. versionadded:: 1.3.0
  163. """
  164. line_type = as_line_type(line_type)
  165. if line_type in (LineType.Separate, LineType.SeparateCode):
  166. # No-op if line_type is not chunked.
  167. return multi_lines
  168. return [dechunk_lines(lines, line_type) for lines in multi_lines]