mpl_util.py 3.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. from __future__ import annotations
  2. from typing import TYPE_CHECKING, cast
  3. import matplotlib.path as mpath
  4. import numpy as np
  5. from contourpy import FillType, LineType
  6. from contourpy.array import codes_from_offsets
  7. if TYPE_CHECKING:
  8. from contourpy._contourpy import FillReturn, LineReturn, LineReturn_Separate
  9. def filled_to_mpl_paths(filled: FillReturn, fill_type: FillType) -> list[mpath.Path]:
  10. if fill_type in (FillType.OuterCode, FillType.ChunkCombinedCode):
  11. paths = [mpath.Path(points, codes) for points, codes in zip(*filled) if points is not None]
  12. elif fill_type in (FillType.OuterOffset, FillType.ChunkCombinedOffset):
  13. paths = [mpath.Path(points, codes_from_offsets(offsets))
  14. for points, offsets in zip(*filled) if points is not None]
  15. elif fill_type == FillType.ChunkCombinedCodeOffset:
  16. paths = []
  17. for points, codes, outer_offsets in zip(*filled):
  18. if points is None:
  19. continue
  20. points = np.split(points, outer_offsets[1:-1])
  21. codes = np.split(codes, outer_offsets[1:-1])
  22. paths += [mpath.Path(p, c) for p, c in zip(points, codes)]
  23. elif fill_type == FillType.ChunkCombinedOffsetOffset:
  24. paths = []
  25. for points, offsets, outer_offsets in zip(*filled):
  26. if points is None:
  27. continue
  28. for i in range(len(outer_offsets)-1):
  29. offs = offsets[outer_offsets[i]:outer_offsets[i+1]+1]
  30. pts = points[offs[0]:offs[-1]]
  31. paths += [mpath.Path(pts, codes_from_offsets(offs - offs[0]))]
  32. else:
  33. raise RuntimeError(f"Conversion of FillType {fill_type} to MPL Paths is not implemented")
  34. return paths
  35. def lines_to_mpl_paths(lines: LineReturn, line_type: LineType) -> list[mpath.Path]:
  36. if line_type == LineType.Separate:
  37. if TYPE_CHECKING:
  38. lines = cast(LineReturn_Separate, lines)
  39. paths = []
  40. for line in lines:
  41. # Drawing as Paths so that they can be closed correctly.
  42. closed = line[0, 0] == line[-1, 0] and line[0, 1] == line[-1, 1]
  43. paths.append(mpath.Path(line, closed=closed))
  44. elif line_type in (LineType.SeparateCode, LineType.ChunkCombinedCode):
  45. paths = [mpath.Path(points, codes) for points, codes in zip(*lines) if points is not None]
  46. elif line_type == LineType.ChunkCombinedOffset:
  47. paths = []
  48. for points, offsets in zip(*lines):
  49. if points is None:
  50. continue
  51. for i in range(len(offsets)-1):
  52. line = points[offsets[i]:offsets[i+1]]
  53. closed = line[0, 0] == line[-1, 0] and line[0, 1] == line[-1, 1]
  54. paths.append(mpath.Path(line, closed=closed))
  55. elif line_type == LineType.ChunkCombinedNan:
  56. paths = []
  57. for points in lines[0]:
  58. if points is None:
  59. continue
  60. nan_offsets = np.nonzero(np.isnan(points[:, 0]))[0]
  61. nan_offsets = np.concatenate([[-1], nan_offsets, [len(points)]])
  62. for s, e in zip(nan_offsets[:-1], nan_offsets[1:]):
  63. line = points[s+1:e]
  64. closed = line[0, 0] == line[-1, 0] and line[0, 1] == line[-1, 1]
  65. paths.append(mpath.Path(line, closed=closed))
  66. else:
  67. raise RuntimeError(f"Conversion of LineType {line_type} to MPL Paths is not implemented")
  68. return paths