draw_nd.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. import numpy as np
  2. def _round_safe(coords):
  3. """Round coords while ensuring successive values are less than 1 apart.
  4. When rounding coordinates for `line_nd`, we want coordinates that are less
  5. than 1 apart (always the case, by design) to remain less than one apart.
  6. However, NumPy rounds values to the nearest *even* integer, so:
  7. >>> np.round([0.5, 1.5, 2.5, 3.5, 4.5])
  8. array([0., 2., 2., 4., 4.])
  9. So, for our application, we detect whether the above case occurs, and use
  10. ``np.floor`` if so. It is sufficient to detect that the first coordinate
  11. falls on 0.5 and that the second coordinate is 1.0 apart, since we assume
  12. by construction that the inter-point distance is less than or equal to 1
  13. and that all successive points are equidistant.
  14. Parameters
  15. ----------
  16. coords : 1D array of float
  17. The coordinates array. We assume that all successive values are
  18. equidistant (``np.all(np.diff(coords) = coords[1] - coords[0])``)
  19. and that this distance is no more than 1
  20. (``np.abs(coords[1] - coords[0]) <= 1``).
  21. Returns
  22. -------
  23. rounded : 1D array of int
  24. The array correctly rounded for an indexing operation, such that no
  25. successive indices will be more than 1 apart.
  26. Examples
  27. --------
  28. >>> coords0 = np.array([0.5, 1.25, 2., 2.75, 3.5])
  29. >>> _round_safe(coords0)
  30. array([0, 1, 2, 3, 4])
  31. >>> coords1 = np.arange(0.5, 8, 1)
  32. >>> coords1
  33. array([0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5])
  34. >>> _round_safe(coords1)
  35. array([0, 1, 2, 3, 4, 5, 6, 7])
  36. """
  37. if len(coords) > 1 and coords[0] % 1 == 0.5 and coords[1] - coords[0] == 1:
  38. _round_function = np.floor
  39. else:
  40. _round_function = np.round
  41. return _round_function(coords).astype(int)
  42. def line_nd(start, stop, *, endpoint=False, integer=True):
  43. """Draw a single-pixel thick line in n dimensions.
  44. The line produced will be ndim-connected. That is, two subsequent
  45. pixels in the line will be either direct or diagonal neighbors in
  46. n dimensions.
  47. Parameters
  48. ----------
  49. start : array-like, shape (N,)
  50. The start coordinates of the line.
  51. stop : array-like, shape (N,)
  52. The end coordinates of the line.
  53. endpoint : bool, optional
  54. Whether to include the endpoint in the returned line. Defaults
  55. to False, which allows for easy drawing of multi-point paths.
  56. integer : bool, optional
  57. Whether to round the coordinates to integer. If True (default),
  58. the returned coordinates can be used to directly index into an
  59. array. `False` could be used for e.g. vector drawing.
  60. Returns
  61. -------
  62. coords : tuple of arrays
  63. The coordinates of points on the line.
  64. Examples
  65. --------
  66. >>> lin = line_nd((1, 1), (5, 2.5), endpoint=False)
  67. >>> lin
  68. (array([1, 2, 3, 4]), array([1, 1, 2, 2]))
  69. >>> im = np.zeros((6, 5), dtype=int)
  70. >>> im[lin] = 1
  71. >>> im
  72. array([[0, 0, 0, 0, 0],
  73. [0, 1, 0, 0, 0],
  74. [0, 1, 0, 0, 0],
  75. [0, 0, 1, 0, 0],
  76. [0, 0, 1, 0, 0],
  77. [0, 0, 0, 0, 0]])
  78. >>> line_nd([2, 1, 1], [5, 5, 2.5], endpoint=True)
  79. (array([2, 3, 4, 4, 5]), array([1, 2, 3, 4, 5]), array([1, 1, 2, 2, 2]))
  80. """
  81. start = np.asarray(start)
  82. stop = np.asarray(stop)
  83. npoints = int(np.ceil(np.max(np.abs(stop - start))))
  84. if endpoint:
  85. npoints += 1
  86. coords = np.linspace(start, stop, num=npoints, endpoint=endpoint).T
  87. if integer:
  88. for dim in range(len(start)):
  89. coords[dim, :] = _round_safe(coords[dim, :])
  90. coords = coords.astype(int)
  91. return tuple(coords)