sketch of hough transform based track finding/fitting

hough.py

#!/usr/bin/env python
import ROOT as r
import sys
import itertools as it
import numpy as np
import pylab as pl
from scipy.special import cotdg
from scipy.sparse import dok_matrix

def hough_radial(points, radius=5):
    points = points.copy()
    points[:,0] -= 0.5
    points[:,1] -= 0.5
    xmax = int(points[:,0].max())+1
    ymax = int(points[:,1].max())+1
    img = dok_matrix((xmax, ymax), dtype=int)

    for px,py,pv in points:
        img[px,py] = pv

    theta_values = np.linspace(-np.pi,np.pi,nbins_theta)
    ct = np.cos(theta_values)
    st = np.sin(theta_values)

    '''
    for x0,y0 in img.keys():
        chunk = img[max(0,x0-radius):x0+radius,max(0,y0-radius):y0+radius]
        for px,py in chunk.keys():
            rvals = (x0+(px-chunk.shape[0]/2)*ct + (y0+(py-chunk.shape[0]/2)*st
    '''
    return img

def hough_lines(points, nmin=3):
    nbins_theta = 180
    theta_values = np.linspace(-np.pi,np.pi,nbins_theta)
    print theta_values.min(), theta_values.max()
    hits = []
    pl.figure()

    ct = np.cos(theta_values)
    st = np.sin(theta_values)

    for p in points:
        x,y,v = p
        #shifts = np.random.normal(loc=0, scale=0.5, size=(100,2))
        shifts = [(0,0)]
        for dx, dy in shifts:
            rvals = (x+dx)*ct + (y+dy)*st
            tvals = theta_values[rvals>0]
            rvals = rvals[rvals>0]
            pl.plot(tvals, rvals)
            if len(rvals):
                hits.extend(list(zip(rvals, tvals)))
    hits = np.array(hits)
    rmax = int(np.ceil(hits[:,0].max()))
    if rmax%2==1: rmax+=1
    h2d = np.histogram2d(hits[:,0], hits[:,1], range=[[0,rmax],[-np.pi,np.pi]], bins=[rmax*10,nbins_theta])
    #print h2d[0]
    print "avg: %.2f, max: %.2f" % (h2d[0].mean(), h2d[0].max())
    #print h2d[0].max()
    print "nmax: ", np.sum(h2d[0]==h2d[0].max())
    print "nover: ", np.sum(h2d[0]>=nmin)
    #pl.matshow(h2d[0])

    argmax = np.unravel_index(h2d[0].argmax(), h2d[0].shape)
    r0 = h2d[1][argmax[0]]
    theta0 = h2d[2][argmax[1]]

    pl.plot([theta0], [r0], 'rx', markersize=20)
    pl.yscale('log')
    pl.ylim(ymin=3e2)

    pl.figure()
    pl.hexbin(hits[:,0], hits[:,1])

    nover = 4 #np.sum(h2d[0]>=nmin)
    candidates = []
    best_candidate = None
    minc2 = np.inf
    for i in xrange(nover):
        loc = np.unravel_index(h2d[0].argmax(), h2d[0].shape)
        h2d[0][loc]=-1
        r0 = h2d[1][loc[0]]
        theta0 = h2d[2][loc[1]]
        candidates.append((r0,theta0))
        c2 = chi2(points, r0, theta0).sum()
        if c2 < minc2:
            best_candidate = (r0, theta0)
            minc2 = c2
    
    pl.plot([r0], [theta0], 'rx', markersize=20)
    print "count at argmax: ", h2d[0][argmax]
    #print r0, theta0
    print 'note: rmax =', rmax
    print "next biggest:", h2d[0][h2d[0]!=h2d[0].max()].max()
    #return r0, theta0

    return best_candidate

def ftrack(points, minpoints=3, minchi2=3):
    points = points.copy()

    while len(points)>minpoints:
        line = hough_lines(points)

        r0, theta0 = line
        c2 = chi2(points, r0, theta0)
        if c2.sum() > minchi2*len(points):
            print "rejecting %d points with (r,theta) = (%.2f, %.2f); red. chi2 = %.2f" % (len(points), r0, theta0, c2.sum()/len(points))
            print "(m=%.2f, b=%.2f)" % trig_to_rect(r0,theta0)
            points = points[c2!=c2.max()]
            line = None
        else:
            print "accepting %d points with (r,theta) = (%.2f, %.2f); red. chi2 = %.2f" % (len(points), r0, theta0, c2.sum()/len(points))
            print "(m=%.2f, b=%.2f)" % trig_to_rect(r0,theta0)
            break

    if line is None: return

    return r0, theta0, points

def hough(points):
    vertices = []
    for p1,p2 in it.combinations(points, r=2):
        x1,y1,v1 = p1
        x2,y2,v1 = p2

        if y1==y2:
            theta = 0
        else:
            theta = np.arctan( 1.*(x1-x2) / (y2-y1) )
        
        r = x1*np.cos(theta) + y1*np.sin(theta)
        if r<0:
            r = -r
            if theta>0: theta-=np.pi
            else: theta += np.pi

        vertices.append((r,theta))
    #print "Got %d vertices from %d points." % (len(vertices), len(points))
    return np.array(vertices)

def hough_rect(points):
    vertices = []
    for p1,p2 in it.combinations(points, r=2):
        x1,y1 = p1
        x2,y2 = p2

        if x1==x2: continue

        m = 1.*(y2-y1) / (x2-x1)
        b = y1 - m*x1
        vertices.append((m,b))
    #print "Got %d vertices from %d points." % (len(vertices), len(points))
    return np.array(vertices)

def find_vertex(vtxs, points):
    minc2 = np.inf
    best_vtx = None
    for v in vtxs:
        c2 = chi2(points, v[0], v[1]).sum()
        if c2<minc2:
            minc2 = c2
            best_vtx = v
    return best_vtx
    #r0 = vtxs[:,0].mean()
    #theta0 = vtxs[:,1].mean()
    #return r0, theta0

def find_vertex_rect(vtxs, median=False):
    if len(vtxs)==0: return
    if median:
        m0 = np.median(vtxs[:,0])
        b0 = np.median(vtxs[:,1])
        return m0, b0
    m0 = vtxs[:,0].mean()
    b0 = vtxs[:,1].mean()
    return m0, b0

def trig_to_rect(r, theta):
    #if np.allclose(theta,0):
    if theta == 0:
        m = np.inf
        b = np.inf
    else:
        m = -1./np.tan(theta)
        b = r/np.sin(theta)
    return m,b

def chi2(points, r, theta):
    #if np.allclose(theta, 0):
    if theta == 0:
        # vertical line, just use xvalues
        dist = (points[:,0] - r)
    else:
        # convert to rectilinear and then
        # use perpendicular distance to line:
        m,b = trig_to_rect(r, theta)
        dist = (m*points[:,0] - points[:,1] + b) / np.sqrt(m**2 + 1)
    #weights = 0.5/points[:,2]
    weights = 0.5
    return (dist/weights)**2

def chi2_rect(points, m, b):
    # use perpendicular distance to line:
    dist = (m*points[:,0] - points[:,1] + b) / np.sqrt(m**2 + 1)
    return (dist/0.5)**2
    #return (points[:,1] - (points[:,0]*m+b))**2

def fit_track(points, maxiter=-1, minpoints=2, minchi2=3, draw=False):
    points = points.copy()

    i = 0
    line = None
    while len(points)>=minpoints:
        i += 1
        if maxiter>0 and i>maxiter: break
        
        vtxs = hough(points)
        
        rvals = vtxs[:,0]
        tvals = vtxs[:,1]
        #wvals = vtxs[:,2]
        rmax = int(np.ceil(rvals.max()))
        if rmax <= 0:
            print "Warning! Rmax is not positive:", rmax
        if rmax%2==1: rmax += 1
        h2d = np.histogram2d(rvals,tvals,range=[[0,rmax],[-np.pi,np.pi]],bins=[rmax*2,360])
        argmax = np.unravel_index(h2d[0].argmax(), h2d[0].shape)
        r0 = h2d[1][argmax[0]]
        theta0 = h2d[2][argmax[1]]
        line = r0,theta0
        if draw:
            pl.figure(0)
            pl.clf()
            pl.matshow(h2d[0], fignum=0)

        if draw:
            pl.figure(i)
            #pl.scatter(vtxs[:,0], vtxs[:,1])
            pl.hexbin(vtxs[:,0], vtxs[:,1])
            pl.xlabel("r")
            pl.ylabel("theta")

        #line = find_vertex(vtxs, points)
        if line is None: return
        r0, theta0 = line
        
        if draw:
            pl.plot([r0], [theta0], 'r+', markersize=20)
            pl.figure(1)
            pl.plot([r0], [theta0], '+', markersize=20)

        c2 = chi2(points, r0, theta0)
        if draw:
            pl.figure(i)
            pl.title('chi2=%.2f'%(c2.sum() / len(points)))

        if c2.sum() > minchi2*len(points):
            points = points[c2!=c2.max()]
            line = None
        else:
            #print "Final red. chi2:", c2.sum()/len(points)
            break

    if line is None:
        return
    
    return r0, theta0, points

def fit_track_rect(points, maxiter=-1, minpoints=2, minchi2=3):
    points = points.copy()

    i = 0
    line = None
    while len(points)>=minpoints:
        i += 1
        if maxiter>0 and i>maxiter: break

        vtxs = hough_rect(points)
        #pl.figure()
        #pl.scatter(vtxs[:,0], vtxs[:,1])
        #pl.plot([vtxs[:,0].mean()], [vtxs[:,1].mean()], 'r+', markersize=20)
        #pl.xlabel("slope")
        #pl.ylabel("offset")
        line = find_vertex_rect(vtxs)
        if not line: return
        m0,b0 = line

        c2 = chi2_rect(points, m0, b0)
        #pl.figure()
        #pl.hist(c2)
        #pl.xlabel("chi^2")
        #pl.ylabel("points")

        if c2.sum() > minchi2*len(points):
            points = points[c2!=c2.max()]
            line = None
        else:
            break

    if not line:
        return

    return m0, b0, points

def find_tracks(t, maxiter=0, rotate=False, get_points=False, minpoints=3, minchi2=3, draw=False):
    if not rotate:
        points = np.array([(x+0.5, y+0.5, v/256.) for x,y,v in zip(t.pix_x, t.pix_y, t.pix_val)])
    else:
        points = np.array([(-y-0.5, x+0.5, v/256.) for x,y,v in zip(t.pix_x, t.pix_y,t.pix_val)])

    tracks = []
    track_points = []
    while len(points) > 3:
        result = ftrack(points, minpoints=minpoints, minchi2=minchi2)

        if not result: break

        r0, theta0, p0 = result

        m0, b0 = trig_to_rect(r0, theta0)

        if m0 == np.inf:
            x0 = p0[:,0].min()
            x1 = p0[:,0].max()
            y0 = p0[:,1].min()
            y1 = p0[:,1].max()

            #assert x0==x1
            if x0!=x1:
                print "Warning! Got infinite slope but delta(X)!=0.", x0,x1,y0,y1

            L = y1-y0
        else:

            x0, x1 = p0[:,0].min(), p0[:,0].max()
            y0 = m0*x0 + b0
            y1 = m0*x1 + b0

            L = np.sqrt( (x1-x0)**2 + (y1-y0)**2 )

        theta = np.arctan(m0)
        c2 = chi2(p0,r0,theta0)
        sum_chi2 = c2.sum()
        if np.isnan(sum_chi2):
            print "WTF! got nan in chi2"
            print "theta0 = ", theta0, "r0 =", r0

        tracks.append((x0,y0,x1,y1,L,theta,len(p0),sum_chi2))
        if get_points:
            track_points.append(p0)

        used = np.array([not p in p0 for p in points])
        points = points[ used ]

    tracks = np.array(tracks, dtype=[('x0',float),('y0',float),('x1',float),('y1',float),('l', float), ('theta', float), ('n', int), ('chi2', float)])

    if get_points:
        return tracks, track_points
    else:
        return tracks


def find_tracks2(t, maxiter=0, rotate=False, get_points=False, minpoints=3, minchi2=3, draw=False):
    if not rotate:
        points = np.array([(x+0.5, y+0.5, v/256.) for x,y,v in zip(t.pix_x, t.pix_y, t.pix_val)])
    else:
        points = np.array([(-y-0.5, x+0.5, v/256.) for x,y,v in zip(t.pix_x, t.pix_y,t.pix_val)])

    tracks = []
    track_points = []
    i = 0
    while len(points) > 2:
        i += 1
        if maxiter>0 and i>maxiter: break

        result = fit_track(points, minpoints=minpoints, minchi2=minchi2, draw=draw)
        if result is None: break

        r0, theta0, p0 = result

        m0, b0 = trig_to_rect(r0, theta0)

        if m0 == np.inf:
            x0 = p0[:,0].min()
            x1 = p0[:,0].max()
            y0 = p0[:,1].min()
            y1 = p0[:,1].max()

            #assert x0==x1
            if x0!=x1:
                print "Warning! Got infinite slope but delta(X)!=0.", x0,x1,y0,y1

            L = y1-y0
        else:

            x0, x1 = p0[:,0].min(), p0[:,0].max()
            y0 = m0*x0 + b0
            y1 = m0*x1 + b0

            L = np.sqrt( (x1-x0)**2 + (y1-y0)**2 )

        theta = np.arctan(m0)
        c2 = chi2(p0,r0,theta0)
        sum_chi2 = c2.sum()
        if np.isnan(sum_chi2):
            print "WTF! got nan in chi2"
            print "theta0 = ", theta0, "r0 =", r0

        tracks.append((x0,y0,x1,y1,L,theta,len(p0),sum_chi2))
        if get_points:
            track_points.append(p0)

        used = np.array([not p in p0 for p in points])
        points = points[ used ]

    tracks = np.array(tracks, dtype=[('x0',float),('y0',float),('x1',float),('y1',float),('l', float), ('theta', float), ('n', int), ('chi2', float)])

    if get_points:
        return tracks, track_points
    else:
        return tracks


def find_tracks_rect(t, maxiter=-1, rotate=False):
    if not rotate:
        points = np.array([(x+0.5, y+0.5) for x,y in zip(t.pix_x, t.pix_y)])
    else:
        points = np.array([(-y-0.5, x+0.5) for x,y in zip(t.pix_x, t.pix_y)])

    tracks = []
    i = 0
    while len(points) > 2 and (maxiter>-1 and i<=maxiter):
        i += 1
        result = fit_track(points)
        if not result: break

        m0, b0, p0 = result
        x0, x1 = p0[:,0].min(), p0[:,0].max()
        y0 = m0*x0 + b0
        y1 = m0*x1 + b0
        L = np.sqrt( (x1-x0)**2 + (y1-y0)**2 )
        theta = np.arctan(m0)
        tracks.append((x0,y0,x1,y1, L, theta, len(p0), chi2_rect(p0,m0,b0).sum()))

        used = np.array([not p in p0 for p in points])
        points = points[ used ]

    tracks = np.array(tracks, dtype=[('x0',float),('y0',float),('x1',float),('y1',float),('l', float), ('theta', float), ('n', int), ('chi2', float)])
    return tracks

if __name__ == "__main__":
    infile = sys.argv[1]

    t0 = r.TChain("event")
    t0.Add(infile)

    '''
    t1 = t0.CloneTree(0)
    t1._pix_hc = r.vector('double')()
    t1.Branch("pix_hc", t1._pix_hc)
    t1._pix_hd = r.vector('double')()
    t1.Branch("pix_hd", t1._pix_hd)

    for evt in t0:
        t1._pix_hc.clear()
        t1._pix_hd.clear()
        for px, py in zip(evt.pix_x, evt.pix_y):
            d = 
    '''