aboutsummaryrefslogtreecommitdiffstats
path: root/fg21sim/utils/draw.py
blob: 25d83c919ddc27f49d0c9081f703d029ce34b43d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
# Copyright (c) 2016 Weitian LI <liweitianux@live.com>
# MIT license


"""
Generic drawers (a.k.a. painters) that draw some commonly used shapes.
"""


import numpy as np
import numba as nb


@nb.jit(nb.types.UniTuple(nb.int64[:], 2)(nb.types.UniTuple(nb.int64, 2),
                                          nb.types.UniTuple(nb.int64, 2),
                                          nb.types.UniTuple(nb.int64, 2)),
        nopython=True)
def ellipse(center, radii, shape):
    """
    Generate coordinates of pixels within the ellipse.

    XXX/NOTE
    --------
    * Cannot figure out why ``nb.optional(nb.types.UniTuple(nb.int64, 2))``
      does NOT work.  Therefore, make ``shape`` as mandatory parameter
      instead of optional.

    Parameters
    ----------
    center : int tuple (r0, c0)
        Center coordinate of the ellipse.
    radii : int tuple (r_radius, c_radius)
        Minor and major semi-axes. ``(r/r_radius)**2 + (c/c_radius)**2 <= 1``.
    shape : int tuple
        Image shape which is used to determine the maximum extent of output
        pixel coordinates.  This is useful for ellipses that exceed the image
        size.  If None, the full extent of the ellipse is used.

    Returns
    -------
    rr, cc : int `~numpy.ndarray`
        Pixel coordinates of the ellipse.
        May be used to directly index into an array, e.g.
        ``img[rr, cc] = 1``.
    """
    # XXX: ``numba`` currently does not support ``numpy.meshgrid``
    nrow, ncol = shape
    r_lim = np.zeros((nrow, ncol))
    for i in range(nrow):
        r_lim[i, :] = np.arange(float(ncol))
    c_lim = np.zeros((nrow, ncol))
    for i in range(ncol):
        c_lim[:, i] = np.arange(float(nrow))

    r_o, c_o = center
    r_r, c_r = radii
    distances = (((r_lim-r_o) / r_r) * ((r_lim-r_o) / r_r) +
                 ((c_lim-c_o) / c_r) * ((c_lim-c_o) / c_r))
    r_idx, c_idx = np.nonzero(distances <= 1.0)
    return (r_idx, c_idx)