#!/Library/Frameworks/Python.framework/Versions/5.1.0/bin/python

import Analysis, re, os, math

from optparse import OptionParser

from collections import defaultdict

from glob import glob

from numpy import array, std, mean, median, Inf, sqrt, diff

from scipy import io, rand

from pylab import plot, quiver, title, quiverkey, figure, ion

from __builtin__ import dir

reload(Analysis)

def gen_plotter(fnames):

sqrt(mean(diff(xs)**2 + diff(ys)**2))/sqrt(2))

if __name__ == "__main__":

parser = OptionParser(usage = "Usage: %prog [opts] files")

# Either we use a mapping file

parser.add_option('-m', '--map', dest = "map",

help = "Mapping file base name ", default = False)

parser.add_option('-2', '--2nd-map', dest = "map2",

help = "Second round map", default = False)

parser.add_option('-S', '--save-offsets', dest="save_map2",

action="store_true", default = False,

help = "Save the offset map for later use")

parser.add_option('-i', '--intra-frame-map', dest="frame_map",

default = 0, type="float",

help = "Calculate an offset map using some fraction of the "

"points in the frame")

# Or we enter the translation manually

parser.add_option('--x-shift', dest = "xshift",

help = "Shift in x between the two channels", default = 0,

type="float")

parser.add_option('-y', '--y-shift', dest = "yshift",

help = "Shift in y between the two channels (default 256 px)",

default = 256, type="float")

# Or we have a single point-pair that should match up, and calculate translation from that

parser.add_option('--x1', dest = "x1", default = 0, type = "float")

parser.add_option('--x2', dest = "x2", default = 0, type = "float")

parser.add_option('--y1', dest = "y1", default = 0, type = "float")

parser.add_option('--y2', dest = "y2", default = 0, type = "float")

parser.add_option('-F', '--frame-by-frame', dest="framewise", default=False,

help = "Process data from each frame separately", action="store_true")

parser.add_option('--flip-y', dest="flip",

help = "Flip the y axis", action="store_true", default = False)

parser.add_option('-f', '--from-channel', dest = "from_chan",

help = "Channel to map from (default 585)", default = "585")

parser.add_option('-t', '--to-channel', dest = "to_chan",

help = "Channel to map onto (default 655)", default = "655")

parser.add_option('-x', '--pixel-size', dest = "pxsize",

help = "Pixel size (default 106.6)", default = 106.667,

type = "float")

parser.add_option('-s', '--spot-std', dest = "max_std",

help = "Maximum allowed inter-frame std dev (default 50)",

default = 50, type="float")

parser.add_option('-b', '--min-brightness', default=8000, type=float,

help = "Minimum brightness for each spot")

parser.add_option('-r', '--max-r', dest = "max_r",

help = "Maximum allowed inter-color colocalization distance",

default = 100, type = "float")

parser.add_option('-R', '--max-2ndround-r', dest="max_R",

default = 20, type="float",

help = "Maximum allowed inter-color colocalization distance "

"in the 2nd round")

parser.add_option('-O', '--max-outlier-r', dest="outlier",

default = 20, type="float",

help="")

parser.add_option('-q', '--quiet', dest = "quiet",

action = "store_true", default = False)

parser.add_option('-o', '--output-file', dest="output", default=None)

parser.add_option('-c', '--center', dest="center", default=False,

help="For measuring colocalization distance, take as distance from"

"the average colocalization", action="store_true")

opts, args = parser.parse_args()

if not opts.map:

if opts.x1 and opts.x2 and opts.y1 and opts.y2:

opts.xshift = opts.x2 - opts.x1

opts.yshift = opts.y2 - opts.y1

else:

mapping = Analysis.loadmapping(opts.map)

if opts.map2:

mapping2 = Analysis.loadmapping(opts.map2)

timepoints = defaultdict(list)

tpmins = [0]*100

tpmaxes = [0]*100

imgname_finder = re.compile('(.*)_spot[0-9]+_[^_]+_xy(_only)?.txt')

spotnum_finder = re.compile('spot([0-9]+)_')

channel_finder = re.compile('spot[0-9]+_([0-9]{3})_xy(_only)?')

#print map(glob, args)

max_timepoint = 0

if len(args) == 1 and '*' in args[0]:

args = glob(args[0])

print "Found this many files", len(args)

DQ_nomatch_x = []

DQ_nomatch_y = []

DQ_nomatch_t = []

DQ_varbig_x = []

DQ_varbig_y = []

DQ_varbig_t = []

DQ_toodim_x = []

DQ_toodim_y = []

DQ_toodim_t = []

xr = []

xl = []

yr = []

yl = []

varxr = []

varxl = []

for fname in args:

try:

spotnum = int(spotnum_finder.findall(fname)[0])

framenum = int(spotnum/100)

imgname = imgname_finder.findall(fname)[0][0].rsplit('/')[-1] + '.tif'

channel = channel_finder.findall(fname)[0][0]

xs, ys, ns, lnums = array([ (float(line.split()[0]),

float(line.split()[1]),

float(line.split()[-1]),

int(lnum))

for lnum, line in enumerate(file(fname))

if line[0].isdigit()]).T

lnums -= 2

tpmin, tpmax = int(min(lnums)), int(max(lnums))

if opts.framewise:

max_timepoint = max(max_timepoint, tpmax)

tpmaxes = array([max_timepoint, 0])

tpmins = array([max_timepoint, 0])

#print tpmaxes

while len(timepoints[channel]) <= tpmax:

timepoints[channel].append([])

if len(xs) == 0:

raise ValueError("Failed on reading somehow")

for x, y, n, lnum in zip(xs, ys, ns, lnums):

lnum = int(lnum)

if n < opts.min_brightness:

DQ_toodim_x.append(x/opts.pxsize)

DQ_toodim_y.append(y/opts.pxsize)

DQ_toodim_t.append(lnum)

elif opts.flip:

timepoints[channel][lnum].append((x,

512*opts.pxsize - y,

n))

else:

timepoints[channel][lnum].append((x, y, n))

else:

if tpmin < tpmins[framenum] or tpmax > tpmaxes[framenum]:

tpmins[framenum] = int(tpmin)

tpmaxes[framenum] = int(tpmax)

if len(xs) < 2: raise ValueError("Bailing on too short")

if std(sqrt(diff(xs)**2 + diff(ys)**2)/sqrt(2)) < opts.max_std:

if mean(ns) < opts.min_brightness:

DQ_toodim_x.append(mean(xs)/opts.pxsize)

DQ_toodim_y.append(mean(ys)/opts.pxsize)

DQ_toodim_t.append(framenum)

continue

while len(timepoints[channel]) <= framenum:

timepoints[channel].append([])

max_timepoint = max(max_timepoint, framenum)

if opts.flip:

timepoints[channel][framenum].append((mean(xs),

512*opts.pxsize - mean(ys),

sqrt(std(xs)**2 + std(ys)**2)))

else:

timepoints[channel][framenum].append((mean(xs),

mean(ys),

sqrt(std(xs)**2 + std(ys)**2)))

else:

DQ_varbig_x.append(mean(xs)/opts.pxsize)

DQ_varbig_y.append(mean(ys)/opts.pxsize)

DQ_varbig_t.append(framenum)

except ValueError as blarg:

print "Not enough data in ", fname, blarg

pass

except:

print "Failed on file", fname, "for some reason..."

raise

print "Max timepoint:", max_timepoint

for channel in timepoints:

print channel, ":", sum(map(len, timepoints[channel]))

try:

goodx = []

goody = []

newx = []

newy = []

diffx = []

diffy = []

framesetnum = []

dists = []

if opts.flip:

opts.to_chan, opts.from_chan = opts.from_chan, opts.to_chan

for i in range(max_timepoint):

if len(timepoints[opts.from_chan]) <= i or \

len(timepoints[opts.to_chan]) <= i:

print "Bailing on timepoint", i

continue

else:

pass

#print "Processing timepoint ", i

for x,y,e in timepoints[opts.from_chan][i]:

if opts.map:

xn, yn = array(mapping(x/opts.pxsize,y/opts.pxsize)) * opts.pxsize

#yn += 10.18

else:

xn = array([((x / opts.pxsize) + opts.xshift) * opts.pxsize])

yn = array([((y / opts.pxsize) + opts.yshift) * opts.pxsize])

if xn == 0 or yn == 0:

continue

bestx = 0

besty = 0

bestd = Inf

beste = 0

for x2, y2, e2 in timepoints[opts.to_chan][i]:

currd = (xn - x2)**2 + (yn - y2)**2

if currd < bestd:

bestd = currd

bestx = x2

besty = y2

beste = e2

dists.append(bestd)

goodx.append(bestx)

goody.append(besty)

newx.append(xn[0])

newy.append(yn[0])

diffx.append(bestx-xn)

diffy.append(besty-yn)

xr.append(x/opts.pxsize)

yr.append(y/opts.pxsize)

varxr.append(e)

xl.append(bestx/opts.pxsize)

yl.append(besty/opts.pxsize)

varxl.append(beste)

framesetnum.append(i)

#print x, y, "\t", bestx, besty, "\t", xn, yn, math.sqrt(bestd)

#print "Now up to a total of ", len(xr)

for x,y, e in timepoints[opts.to_chan][i]:

if x not in goodx and y not in goody:

DQ_nomatch_x.append(x/opts.pxsize)

DQ_nomatch_y.append(y/opts.pxsize)

DQ_nomatch_t.append(i)

if opts.center:

centerx = mean(diffx)

centery = mean(diffy)

else:

centerx = 0

centery = 0

print len(diffx)

assert len(varxl) == len(xl)

for i in reversed(range(len(diffx))):

# Count from the back so removing elements doesn't shift numbering

if dists[i] is Inf or \

(diffx[i] - centerx)**2 + (diffy[i] - centery)**2 >= opts.max_r**2:

DQ_nomatch_x.append(xr.pop(i))

DQ_nomatch_y.append(yr.pop(i))

fn = framesetnum.pop(i)

DQ_nomatch_t.append(fn)

goodx.pop(i)

goody.pop(i)

newx.pop(i)

newy.pop(i)

diffx.pop(i)

diffy.pop(i)

# xr and yr already popped

varxr.pop(i)

DQ_nomatch_x.append(xl.pop(i))

DQ_nomatch_y.append(yl.pop(i))

DQ_nomatch_t.append(fn)

varxl.pop(i)

# Framesetnum already popped

continue

assert len(xl) == len(varxl)

print "New: ", len(diffx)

diffx = array(diffx)

diffy = array(diffy)

newx = array(newx)

newy = array(newy)

goodx = array(goodx)

goody = array(goody)

xl = array(xl)

yl = array(yl)

xr = array(xr)

yr = array(yr)

varxl = array(varxl)

varxr = array(varxr)

framesetnum = array(framesetnum)

print "Colocalized ", len(goodx), " spots at an average error of ",

print median(sqrt(diffx**2 + diffy**2))

if opts.map2:

print "Applying fiducials"

xds, yds = mapping2(newx/opts.pxsize, newy/opts.pxsize)

newx += xds

newy += yds

diffx = goodx - newx

diffy = goody - newy

print diffx[0], diffy[0]

sel2 = (abs(xds - centerx) < opts.max_R ) & (abs(yds-centery) < opts.max_R) \

& (diffx < opts.max_R) & (diffy < opts.max_R)

diffx = diffx[sel2]

diffy = diffy[sel2]

goodx = goodx[sel2]

goody = goody[sel2]

xl = xl[sel2]

yl = yl[sel2]

xr = xr[sel2]

yr = yr[sel2]

newx = newx[sel2]

newy = newy[sel2]

varxl = varxl[sel2]

varxr = varxr[sel2]

framesetnum = framesetnum[sel2]

print "In the second round, those colocalized at an error of",

print median(sqrt(diffx**2 + diffy**2))

elif opts.frame_map:

new_goodx = []

new_goody = []

new_diffx = []

new_diffy = []

for i in range(1):#range(max(framesetnum)):

#splitter = (framesetnum == 0)

frame_select = array([True]*len(framesetnum))

splitter = (rand(len(xl)) < opts.frame_map)

#splitter = array([True]*len(framesetnum))

num_spots_used = sum(frame_select & splitter)

print "On Frame %d, using %d spots for regression" % (i, num_spots_used)

if not num_spots_used: continue

diffx_fids = diffx[frame_select & splitter]

diffy_fids = diffy[frame_select & splitter]

newx_fids = newx [frame_select & splitter]

newy_fids = newy [frame_select & splitter]

goodx_fids = goodx[frame_select & splitter]

goody_fids = goody[frame_select & splitter]

xl_fids = xl [frame_select & splitter]

yl_fids = yl [frame_select & splitter]

xr_fids = xr [frame_select & splitter]

yr_fids = yr [frame_select & splitter]

framesetnum_fids = framesetnum[frame_select & splitter]

diffx_exp = diffx[frame_select & ~splitter]

diffy_exp = diffy[frame_select & ~splitter]

newx_exp = newx [frame_select & ~splitter]

newy_exp = newy [frame_select & ~splitter]

goodx_exp = goodx[frame_select & ~splitter]

goody_exp = goody[frame_select & ~splitter]

xl_exp = xl [frame_select & ~splitter]

yl_exp = yl [frame_select & ~splitter]

xr_exp = xr [frame_select & ~splitter]

yr_exp = yr [frame_select & ~splitter]

framesetnum_exp = framesetnum[frame_select & ~splitter]

if opts.save_map2:

mapping2 = Analysis.makeregression(diffx_fids, diffy_fids,

newx_fids, newy_fids, order=1,

savefile = 'offsets'

+ os.path.splitext(os.path.basename(imgname))[0],)

else:

mapping2 = Analysis.makeregression(diffx_fids, diffy_fids,

newx_fids, newy_fids, order=1)

dx, dy = mapping2(newx_exp, newy_exp)

newx_exp += dx

newy_exp += dy

new_goodx.extend(goodx_exp)

new_goody.extend(goody_exp)

new_diffx.extend(goodx_exp - newx_exp)

new_diffy.extend(goody_exp - newy_exp)

goodx = array(new_goodx)

goody = array(new_goody)

diffx = array(new_diffx)

diffy = array(new_diffy)

sel2 = (dx < opts.max_R ) & (dy < opts.max_R) \

& (diffx < opts.max_R) & (diffy < opts.max_R)

diffx = diffx[sel2]

diffy = diffy[sel2]

goodx = goodx[sel2]

goody = goody[sel2]

#xl = xl[sel2]

#yl = yl[sel2]

#xr = xr[sel2]

#yr = yr[sel2]

#newx = newx[sel2]

#newy = newy[sel2]

print "In the second round, those colocalized at an error of",

print median(sqrt(diffx**2 + diffy**2))

print "While keeping ", len(goodx), " spots"

else:

if opts.save_map2 \

or raw_input('Save Second Pass? y/[n] ').lower() == 'y':

mapping2 = Analysis.makeregression(diffx, diffy, xl, yl,

savefile= 'offsets'

+ os.path.splitext(os.path.basename(imgname))[0],

order = int(max(2,sqrt(len(goody)/10))))

# for x1,y1, x2,y2, e in zip(xr, yr, xl, yl, varxl):

# if opts.map:

# print "#", x1, y1, "\t", x2, y2, "\t", array(mapping(x1, y1)).T, "\t\t", math.sqrt(e)

except IndexError:

print "Not enough data!"

finally:

# Remove outliers from plots and output. Presumably, this could have a

# different value from max_R. Otherwise, I'm not sure what the

# difference is, but Ahmet did ask for it.

sel2 = (diffx**2 + diffy**2) < opts.outlier**2

sel2 = sel2.squeeze()

xl = xl[sel2]

yl = yl[sel2]

xr = xr[sel2]

yr = yr[sel2]

varxl = varxl[sel2]

varxr = varxr[sel2]

newx = newx[sel2]

newy = newy[sel2]

goodx = goodx[sel2]

goody = goody[sel2]

diffx = diffx[sel2]

diffy = diffy[sel2]

try:

print "Saving to ", os.path.dirname(fname)+'.mat'

output_dict = {}

for var in ('xl', 'yl', 'xr', 'yr', 'varxl', 'varxr', 'newx', 'newy',

'diffx', 'diffy',

'framesetnum', 'tpmins', 'tpmaxes', 'imgname',

'DQ_varbig_x','DQ_varbig_y', 'DQ_varbig_t',

'DQ_nomatch_x', 'DQ_nomatch_y', 'DQ_nomatch_t',

'DQ_toodim_x', 'DQ_toodim_y', 'DQ_toodim_t'):

if var in dir():

exec 'output_dict["%s"] = %s' % (var, var)

if opts.frame_map:

for var in ('xl_fids', 'yl_fids', 'xr_fids', 'yr_fids',

'xl_exp', 'yl_exp', 'xr_exp', 'yr_exp',

'framesetnum_fids', 'framesetnum_exp'):

if var in dir():

exec 'output_dict["%s"] = %s' % (var, var)

output_dict['framebyframe'] = opts.framewise

output_dict['mapname'] = opts.map

output_dict['mapname2'] = opts.map2

io.savemat(opts.output or os.path.dirname(fname) + '.mat', output_dict)

except IOError:

print "Couldn't save a .mat file. Probably a filesystem issue"

if not opts.quiet:

ion()

figure()

Q = quiver(goodx, goody, diffx, diffy, angles='xy', minshaft=2, units='x')

scale = math.sqrt(median(array(diffx)**2 + array(diffy)**2))

title('%s with mapping %s and error < %d' %

(os.path.dirname(fname) ,str(opts.map or opts.yshift), opts.max_r))

print scale

if scale < 50:

quiverkey(Q, .1, .1, 10, '$10 nm$', color="blue")

else:

quiverkey(Q, .1, .1, 1000, r'$1\mu m$', color="blue")

figure()

plot(diffx, diffy, 'ro')

title(os.path.dirname(fname))

for channel in timepoints:

num = sum(len(timepoint) for timepoint in timepoints[channel])

print channel, num