How can I draw lines into numpy arrays?

I stumbled on this question while looking for a solution, and the provided answer solves it quite well. However, it didn't really suit my purposes, for which I needed a "tensorizable" solution (i.e. implemented in numpy without explicit loops), and possibly with a linewidth option. I ended up implementing my own version, and since in the end it's also quite faster than line_aa, I thought I could share it.

It comes in two flavors, with and without linewidth. Actually the former is not a generalization of the latter, and neither perfectly agrees with line_aa, but for my purposes they're just fine and on plots they look okay.

def naive_line(r0, c0, r1, c1):
    # The algorithm below works fine if c1 >= c0 and c1-c0 >= abs(r1-r0).
    # If either of these cases are violated, do some switches.
    if abs(c1-c0) < abs(r1-r0):
        # Switch x and y, and switch again when returning.
        xx, yy, val = naive_line(c0, r0, c1, r1)
        return (yy, xx, val)

    # At this point we know that the distance in columns (x) is greater
    # than that in rows (y). Possibly one more switch if c0 > c1.
    if c0 > c1:
        return naive_line(r1, c1, r0, c0)

    # We write y as a function of x, because the slope is always <= 1
    # (in absolute value)
    x = np.arange(c0, c1+1, dtype=float)
    y = x * (r1-r0) / (c1-c0) + (c1*r0-c0*r1) / (c1-c0)

    valbot = np.floor(y)-y+1
    valtop = y-np.floor(y)

    return (np.concatenate((np.floor(y), np.floor(y)+1)).astype(int), np.concatenate((x,x)).astype(int),
            np.concatenate((valbot, valtop)))

I called this "naive" because it is quite similar to the naive implementation in Wikipedia, but with some anti-aliasing, although admittedly not perfect (e.g. makes very thin diagonals).

The weighted version gives much thicker line more pronounced anti-aliasing.

def trapez(y,y0,w):
    return np.clip(np.minimum(y+1+w/2-y0, -y+1+w/2+y0),0,1)

def weighted_line(r0, c0, r1, c1, w, rmin=0, rmax=np.inf):
    # The algorithm below works fine if c1 >= c0 and c1-c0 >= abs(r1-r0).
    # If either of these cases are violated, do some switches.
    if abs(c1-c0) < abs(r1-r0):
        # Switch x and y, and switch again when returning.
        xx, yy, val = weighted_line(c0, r0, c1, r1, w, rmin=rmin, rmax=rmax)
        return (yy, xx, val)

    # At this point we know that the distance in columns (x) is greater
    # than that in rows (y). Possibly one more switch if c0 > c1.
    if c0 > c1:
        return weighted_line(r1, c1, r0, c0, w, rmin=rmin, rmax=rmax)

    # The following is now always < 1 in abs
    slope = (r1-r0) / (c1-c0)

    # Adjust weight by the slope
    w *= np.sqrt(1+np.abs(slope)) / 2

    # We write y as a function of x, because the slope is always <= 1
    # (in absolute value)
    x = np.arange(c0, c1+1, dtype=float)
    y = x * slope + (c1*r0-c0*r1) / (c1-c0)

    # Now instead of 2 values for y, we have 2*np.ceil(w/2).
    # All values are 1 except the upmost and bottommost.
    thickness = np.ceil(w/2)
    yy = (np.floor(y).reshape(-1,1) + np.arange(-thickness-1,thickness+2).reshape(1,-1))
    xx = np.repeat(x, yy.shape[1])
    vals = trapez(yy, y.reshape(-1,1), w).flatten()

    yy = yy.flatten()

    # Exclude useless parts and those outside of the interval
    # to avoid parts outside of the picture
    mask = np.logical_and.reduce((yy >= rmin, yy < rmax, vals > 0))

    return (yy[mask].astype(int), xx[mask].astype(int), vals[mask])

The weight adjustment is admittedly quite arbitrary, so anybody can adjust that to their tastes. The rmin and rmax are now needed to avoid pixels outside of the picture. A comparison:

As you can see, even with w=1, weighted_line is a bit thicker, but in a kind of homogeneous way; similarly, naive_line is homogeneously slightly thinner.

Final note about benchmarking: on my machine, running %timeit f(1,1,100,240) for the various functions (w=1 for weighted_line) resulted in a time of 90 µs for line_aa, 84 µs for weighted_line (although the time of course increases with the weight) and 18 µs for naive_line. Again for comparison, reimplementing line_aa in pure Python (instead of Cython as in the package) took 350 µs.

Thanks to Joe Kington for the answer! I was looking for skimage.draw.line_aa.

import scipy.misc
import numpy as np
from skimage.draw import line_aa
img = np.zeros((10, 10), dtype=np.uint8)
rr, cc, val = line_aa(1, 1, 8, 4)
img[rr, cc] = val * 255
scipy.misc.imsave("out.png", img)