In [1]:

    

import os


    

In [27]:

    

def insert_theta(line,th):
    p = " ".join(line.split()[:4])
    q = " ".join(line.split()[5:])
    return p+" "+str(th)+" "+q


    

In [ ]:

    

import numpy as np
import os
cos, sin = np.cos, np.sin
twopi = np.pi*2.0

def theta_correct(fname,isnap=None,inplace=False):
    # readjust all theta values in a file
    
    foutname = fname+"_thcorr"
    
    # Count num blocks
    Nblock = 0
    fin = open(fname, "r")
    for line in fin.readlines():
        if line == "\n": Nblock+=1
    fin.seek(0)

    ln = fin.readline().split("|")
    nrod = 0
    edge = 0.
    Nx = 0
    for s in ln:
        if "boxEdge" in s:
            edge = float(s.split()[1])
        if "nObj" in s:
            nrod = int(s.split()[1])
        if "cellNx" in s:
            Nx = int(s.split()[1])
    fin.seek(0)
    
    thgrid = np.zeros(shape=(Nx,Nx))
    cell2grid = {}
    for i in range(Nx*Nx):
        xi = i//Nx
        yi = i%Nx
        cell2grid.update{i:[xi,yi]}
    
    # dict of corrected refs so that in subsequent blocks
    # we only need to compare it with the ref
    threfs = {}
    neighbors = {}

    snaps = []
    if isnap:
        snaps = [isnap]
        foutname = foutname+"_"+str(isnap)
    else:
        for s in range(0,Nblock):
            snaps.append(s)
            
    fout = open(foutname, 'w')
    l = fin.readline()
    if l[0].isalpha():
        fout.write(l)
    else:
        fin.seek(0)
    
    xs,ys,thetas = [],[],[]
    rids, cellids = [], []
    blocklines = []
    tmpthetas = []
    cntsnap = 0
    lblstring = None
    
    #
    # First build up threfs based on first image
    #
    for line in fin.readlines():
        if cntsnap not in snaps:
            if line == "\n": cntsnap+=1
            continue
        else:
            if line == "\n":
                # Done the block
                for x,y,th,r,c in zip(xs,ys,thetas,rids,cellids):
                    # th is initially in range [0,2pi]
                    # arctan2 outputs in range [-pi,pi]
                    # difference between arctan and arctan2 is that
                    # arctan only "folds back" the range [-pi/2,pi/2]
                    if th > np.pi: th = -twopi + th # th: [-pi,pi]
                    xi,yi = cell2grid[c]
                    if not thgrid[xi,yi]:
                        # First rod in cell, define direction
                        # get it a 'safe' distance from modpoint for
                        # angle comparisons later
                        if th > np.pi*0.5: th -= np.pi
                        if th < -np.pi*0.5: th += np.pi
                        # now th is [-pi/2,pi/2]
                        thgrid[xi,yi] = th
                    else:
                        thref = thgrid_refs[xi,yi]
                        if (th > thref) and ((th - thref) > 0.25*twopi): th -= 0.5*twopi
                        if (th < thref) and ((thref - th) > 0.25*twopi): th += 0.5*twopi

                    # return theta to [0,2pi] range
                    th += np.pi
                    threfs.update({r: th})
                    tmpthetas.append(th)

                # Reset arrays
                xs = []
                ys = []
                thetas = []
                break
                
            blocklines.append(line)
            spt = [float(x) for x in line.split()]
            xs.append(spt[2])
            ys.append(spt[3])
            thetas.append(spt[4])
            rids.append(int(spt[0]))
            cellids.append(int(spt[1]))

    # If given isnap, we are done
    if isnap:
        ith = 0
        for l in blocklines:
            fout.write(insert_theta(l,tmpthetas[ith]))
            ith+=1
        fout.close()
        return

    xs,ys,thetas = [],[],[]
    blocklines = []

    #
    # Continue on updating rest of the file
    #
    for line in fin.readlines():
        if line.startswith("label") or line == "\n":
            fout.write(line)
            continue

        spt = [float(x) for x in line.split()]
        th = spt[4]
        if th > np.pi: th = -twopi+th
        r = int(spt[0])
        th_ = threfs[r]
        if th_ > np.pi: th_ = -twopi+th_
        if (th > th_) and ((th - th_) > 0.25*twopi): th -= 0.5*twopi
        if (th < th_) and ((th_ - th) > 0.25*twopi): th += 0.5*twopi
        
        # return theta to [0,2pi] range
        th += np.pi
        threfs[r] = th
        
        fout.write(insert_theta(line,th))

    fin.close()
    fout.close()
    print "Done correcting"

    # if inplace = True and isnap=None, overwrite fname
    if inplace and not isnap:
        print "Overwriting",fname
        os.rename(foutname,fname)
    else:
        print "Writing to",foutname