dataio.py

import numpy as np
import matplotlib.pyplot as plt
import collections as co
import cairo
import math
import pdb
import copy
from collections import deque
import os
import scipy.io as sio
import scipy.misc as scm
import pickle
#Custom packages
import primitives as pm
import geometry as gm
import physics as phy
import os
from os import path as osp

class DataSaver:	
	def __init__(self, rootPath='/work5/pulkitag/projPhysics', numBalls=1,
							 mnBallSz=15, mxBallSz=35,
							 mnSeqLen=40, mxSeqLen=100, 
							 mnForce=1e+3, mxForce=1e+6, wThick=30,
							 isRect=True, wTheta=30, mxWLen=600, mnWLen=200,
							 arenaSz=667, oppForce=False,
							 svPrefix=None, randSeed=None, verbose=0, **kwargs):
		'''
			isRect  : If the walls need to be rectangular
			wTheta  : If the walls are NOT rectangular then at what angles should they be present.
			mxWLen  : Maximum length of the walls
			mnWLen  : Minimum length of the walls
			arenaSz : Size of the arena 
			svPrefix: Prefix in the file names for saving the data 
		'''
		#print (rootPath)
		#The name of the experiment. 
		self.expStr_   = 'aSz%d_wLen%d-%d_nb%d_bSz%d-%d_f%.2e-%.2e_sLen%d-%d_wTh%d' % (arenaSz,
										  mnWLen, mxWLen, numBalls, mnBallSz, mxBallSz, mnForce, mxForce,
										  mnSeqLen, mxSeqLen, wThick)
		if svPrefix is not None:
			self.expStr_ = svPrefix + '-' + self.expStr_

		if not isRect:
			if isinstance(wTheta, list):
				thetaStr = '_wTheta'.join('%d-' % th for th in wTheta)
				thetaStr = 	thetaStr[0:-1]
			else:
				thetaStr = '_wTheta%d' % wTheta
			self.expStr_ = self.expStr_ + thetaStr
		if oppForce:
			self.expStr_   = self.expStr_ + '_oppFrc'	

		#pdb.set_trace()
		#Setup directories.
		self.dirName_  = os.path.join(rootPath, self.expStr_)
		if not os.path.exists(self.dirName_):
			os.makedirs(self.dirName_)
		self.seqDir_   = os.path.join(self.dirName_, 'seq%06d')
		self.mnSeqLen_   = mnSeqLen
		self.mxSeqLen_   = mxSeqLen
		self.imFile_    = 'im%06d.jpg'
		self.dataFile_  = 'data.mat'
		self.worldFile_ = 'world.pkl' #Saves the world. 
		#Setup variables.  
		self.numBalls_ = numBalls
		self.bmn_      = mnBallSz
		self.bmx_      = mxBallSz
		self.fmn_      = mnForce
		self.fmx_      = mxForce
		self.wlmx_     = mxWLen
		self.wlmn_     = mnWLen
		self.xSz_      = arenaSz
		self.ySz_      = arenaSz
		self.wth_      = wThick
		self.isRect_   = isRect
		self.oppForce_ = oppForce
		self.verbose_  = verbose
		if not isinstance(wTheta, list):
			wTheta = [wTheta]
		self.wTheta_   = wTheta	
		if randSeed is None:
			self.rand_ = np.random.RandomState()
		else:
			self.rand_ = np.random.RandomState(randSeed)
	print ('DATAIO SETUP DONE')

	def save(self, numSeq=10):
		for i in range(numSeq):
			print i
			seqLen = int(self.mnSeqLen_ + self.rand_.rand() * (self.mxSeqLen_ - self.mnSeqLen_))
			self.seqLen_ = seqLen
			seqDir = self.seqDir_ % i
			if not os.path.exists(seqDir):
				os.makedirs(seqDir)
			dataFile  = os.path.join(seqDir, self.dataFile_)
			imFile    = os.path.join(seqDir, self.imFile_) 
			worldFile = os.path.join(seqDir, self.worldFile_)
			self.save_sequence(dataFile, imFile, worldFile)

	def save_sequence(self, dataFile, imFile, worldFile):
		model, f, ballPos, walls = self.generate_model()
		force    = np.zeros((2 * self.numBalls_, self.seqLen_)).astype(np.float32)
		position = np.zeros((2 * self.numBalls_, self.seqLen_)).astype(np.float32)

		#Collect all the objects in the worlds
		#objs = {}
		#for name in self.world_.get_object_names():
		#	objs[name] = self.world_.get_object(name)
		#pdb.set_trace()
		pickle.dump({'force': f, 'ballPos': ballPos, 'walls': self.pts}, open(worldFile,'w'))

		for b in range(self.numBalls_):
			fb = f[b]
			st, en = 2*b, 2*b + 1
			force[st,0], force[en,0] = fb.x(), fb.y()
			print fb 
		for i in range(self.seqLen_):
			model.step()
			im = model.generate_image()
			svImFile = imFile % i
			scm.imsave(svImFile, im)				
			for j in range(self.numBalls_):
				ballName = 'ball-%d' % j
				ball     = model.get_object(ballName)
				pos      = ball.get_position()
				position[2*j,   i] = pos.x()
				position[2*j+1, i] = pos.y()
		sio.savemat(dataFile, {'force': force, 'position': position})	

	def fetch(self, cropSz=None, procSz=None):
		seqLen = int(self.mnSeqLen_ + self.rand_.rand() * (self.mxSeqLen_ - self.mnSeqLen_))
		self.seqLen_ = seqLen
		model, f, ballPos, walls = self.generate_model()
		force    = np.zeros((2 * self.numBalls_, self.seqLen_)).astype(np.float32)
		position = np.zeros((2 * self.numBalls_, self.seqLen_)).astype(np.float32)
		velocity = np.zeros((2 * self.numBalls_, self.seqLen_)).astype(np.float32)
		imList  = []
		imBalls = []
		for b in range(self.numBalls_):
			imBalls.append([])
			fb = f[b]
			st, en = 2*b, 2*b + 1
			force[st,0], force[en,0] = fb.x(), fb.y()
		#Previous position
		pPos = np.nan * np.zeros((self.numBalls_,2))
		for i in range(self.seqLen_):
			model.step()
			im = model.generate_image()
			vx, vy = None, None
			for j in range(self.numBalls_):
				ballName = 'ball-%d' % j
				ball     = model.get_object(ballName)
				pos      = ball.get_position()
				position[2*j,   i] = pos.x()
				position[2*j+1, i] = pos.y()
				#Speed should not be predicted, instead we should just predict
				#delta in position. The difference between the two is critical 
				#due to collisions. 
				if not np.isnan(pPos[j][0]):
					vx = pos.x() - pPos[j][0]
					vy = pos.y() - pPos[j][1]
				pPos[j][0], pPos[j][1] = pos.x(), pos.y()
				xMid, yMid = round(pos.x()), round(pos.y())
				if cropSz is not None:
					imBall = 255 * np.ones((cropSz, cropSz,3)).astype(np.uint8)
					#Cropping coordinates in the original image
					x1, x2 = max(0, xMid - cropSz/2.0), min(self.xSz_, xMid + cropSz/2.0)
					y1, y2 = max(0, yMid - cropSz/2.0), min(self.ySz_, yMid + cropSz/2.0)
					#Coordinates in the cropped image centerd at the ball
					imX1 = int(round(cropSz/2.0  - (xMid - x1)))
					imX2 = int(round(cropSz/2.0  + (x2 - xMid)))
					imY1 = int(round(cropSz/2.0  - (yMid - y1)))
					imY2 = int(round(cropSz/2.0  + (y2 - yMid)))
					x1, x2 = int(round(x1)), int(round(x2))
					y1, y2 = int(round(y1)), int(round(y2))
					imBall[imY1:imY2,imX1:imX2,:] = im[y1:y2, x1:x2,0:3]
					position[2*j, i] = position[2*j, i] - x1 + imX1 
					position[2*j+1, i] = position[2*j+1, i] - y1 + imY1 

					if procSz is not None:
						posScale  = float(procSz)/float(cropSz)
						imBall = scm.imresize(imBall, (procSz, procSz))
						position[2*j, i]   = position[2*j, i] * posScale
						position[2*j+1, i] = position[2*j+1, i] * posScale
						if vx is not None:
							velocity[2*j, i-1]   = vx * posScale
							velocity[2*j+1, i-1] = vy * posScale
					imBalls[j].append(imBall)
			imList.append(im)
		if cropSz is None:
			return imList
		else:
			return imBalls, force[:, 0:self.seqLen_],\
						 velocity[:, 0:self.seqLen_],\
						 position[:, 0:self.seqLen_] 
		
	def _generate_model(self, returnPos=False):
		#get the coordinates of the top point
		#create the world
		self.world_ = pm.World(xSz=self.xSz_, ySz=self.ySz_)
		#add the walls
		walls = self.add_walls()
		#add the balls
		ballpos = self.add_balls()
		#create physics simulation
		model = pm.Dynamics(self.world_)
		if returnPos:
			return model, ballpos, self.pts, walls
		else:
			return model

	def generate_model(self):
		model, ballpos, _, walls = self._generate_model(returnPos=True)
		#apply initial forces and return the result. 	
		model, fs =  self.apply_force(model)	
		return model, fs, ballpos, walls

	#this is mostly due to legacy reasons. 
	def add_rectangular_walls(self, fColor=pm.Color(1.0, 0.0, 0.0)):
		#define the extents within which walls can be put. 
		hlen   = np.floor(self.wlmn_ + self.rand_.rand() * (self.wlmx_ - self.wlmn_))
		vlen   = np.floor(self.wlmn_ + self.rand_.rand() * (self.wlmx_ - self.wlmn_))
		topxmx = self.xSz_ - (hlen + self.wth_)
		topymx = self.ySz_ - (vlen + self.wth_)
		xleft = np.floor(self.rand_.rand() * topxmx)
		ytop  = np.floor(self.rand_.rand() * topymx)	
		walls  = self._create_walls(xleft, ytop, (0, 90, 180), (hlen - self.wth_, vlen, hlen - self.wth_), 
															 fColor=fColor)
	
		#define the walls
		#wallhordef = pm.walldef(sz=gm.Point(hlen, self.wth_), fColor=fColor)
		#wallverdef = pm.walldef(sz=gm.Point(self.wth_, vlen), fColor=fColor)
		#self.world_.add_object(wallverdef, initpos=gm.Point(xleft, ytop))
		#self.world_.add_object(wallverdef, initpos=gm.Point(xleft + hlen - self.wth_, ytop))
		#self.world_.add_object(wallhordef, initpos=gm.Point(xleft, ytop))
		#self.world_.add_object(wallhordef, initpos=gm.Point(xleft, ytop + vlen))
		#self.pts = [gm.Point(xleft, ytop)]
		#self.whl_, self.wvl_ = hlen, vlen
		return walls

	##
	def sample_walls(self):
		#for adding diagonal walls
		#1. estimate the x and y extents of the wall. 
		#2. find the appropriate starting position based on that
		#sample the theta
		perm = self.rand_.permutation(len(self.wTheta_))
		wtheta = self.wTheta_[perm[0]]
		rad  = (wtheta * np.pi)/180.0
		hlen   = self.wlmn_ + self.rand_.rand() * (self.wlmx_ - self.wlmn_)
		if wtheta == 90:
			xlen = hlen
			ylen = hlen
		else:
			xlen = hlen * np.cos(rad)
			ylen = hlen * np.sin(rad)
		xextent = 2 * xlen +  2 * self.wth_
		yextent = 2 * ylen +  2 * self.wth_
		xleftmin = self.wth_  
		xleftmax = self.xSz_ - xextent
		yleftmin  = ylen + self.wth_
		if wtheta == 90:
			yleftmax = self.ySz_ - self.wth_
		else:
			yleftmax  = self.ySz_ - (ylen + self.wth_)
		#keep sampling until the appropriate size has been found. 
		if xleftmin <= 0 or yleftmin <=0:
			return self.sample_walls()
		if xleftmax < xleftmin or yleftmax < yleftmin:
			return self.sample_walls()
		xleft    = xleftmin + np.floor(self.rand_.rand() * (xleftmax - xleftmin))
		yleft    = yleftmin  + np.floor(self.rand_.rand() * (yleftmax - yleftmin))	
		return xleft, yleft, wtheta, hlen	

	##
	def _create_walls(self, xleft, yleft, thetas, wlens, fColor):
		theta1, theta2, theta3 = thetas
		wlen1,  wlen2,  wlen3  = wlens
		pt1      = gm.Point(xleft, yleft)
		dir1     = gm.theta2dir(theta1)
		pt2      = pt1 + (wlen1 * dir1)
		dir2     = gm.theta2dir(theta2)
		pt3      = pt2 + (wlen2 * dir2)
		dir3     = gm.theta2dir(theta3)
		pt4      = pt3 + (wlen3 * dir3)
		pts      = [pt1, pt2, pt3, pt4]
		if self.verbose_ > 0:
			print ("points: ", pt1, pt2, pt3, pt4)
		walls    = pm.create_cage(pts, wThick = self.wth_, fColor=fColor)	
		#get the lines within which the balls need to be added. 
		self.pts    = pts
		self.lines_ = []
		for w in walls:
			self.world_.add_object(w)
		for i in range(len(pts)):
			self.lines_.append(gm.Line(pts[i], pts[np.mod(i+1, len(pts))]))		
		return walls

	def add_walls(self, fColor=pm.Color(1.0, 0.0, 0.0)):
		if self.isRect_:
			return self.add_rectangular_walls(fColor=fColor)
		xleft, yleft, wtheta, hlen = self.sample_walls()
		walls = self._create_walls(xleft, yleft, (-wtheta, wtheta, 180-wtheta),
							 (hlen, hlen, hlen), fColor=fColor)
		return walls		


	def find_point_within_lines(self, mindist):
		'''
			find a point within the lines which is atleast mindist
			from all the boundaries. 
		'''		
		x = int(np.round(self.pts[0].x() + self.rand_.rand()*(self.pts[2].x() - self.pts[0].x())))
		y = int(np.round(self.pts[1].y() + self.rand_.rand()*(self.pts[3].y() - self.pts[1].y())))
		pt = gm.Point(x,y)
		isinside = True
		dist = []
		for (i,l) in enumerate(self.lines_):
			#note we are finding the signed distance
			dist.append(l.distance_to_point(pt))
			if dist[i] <= mindist:
				isinside=False
		md = min(dist)
		return pt, isinside, md
					
	#generates and adds the required number of balls. 
	def add_balls(self):
		#generate ball definitions
		allr, allpos = [], []
		for i in range(self.numBalls_):
			placeflag = True
			while placeflag:
					#randomly sample the radius of the ball
					r    = int(np.floor(self.bmn_ + self.rand_.rand() * (self.bmx_ - self.bmn_))) 
					bdef = pm.BallDef(radius=r, fColor=pm.Color(0.5, 0.5, 0.5))
					#find a position to keep the ball
					'''
					if self.isrect_:
						xleft, ytop = self.pts[0].x_asint(), self.pts[0].y_asint()
						#xmn  = xleft + 2 * r + self.wth_
						#ymn  = ytop  + 2 * r + self.wth_
						#xmx  = xleft + self.whl_ - self.wth_ - 2 * r
						#ymx  = ytop  + self.wvl_ - self.wth_ - 2 * r
						xmn  = xleft + r + self.wth_ + 2 #give some margin
						ymn  = ytop  + r + self.wth_ + 2
						xmx  = xleft + self.whl_ - self.wth_ - r - 2
						ymx  = ytop  + self.wvl_ - self.wth_ - r - 2
						xloc = int(np.floor(xmn + (xmx - xmn) * self.rand_.rand()))
						yloc = int(np.floor(ymn + (ymx - ymn) * self.rand_.rand()))
					else:
					'''
					findflag = True
					count    = 0
					while findflag:
						pt, isvalid, md = self.find_point_within_lines(r + self.wth_ + 2) #2 is safety margin	
						count += 1
						if isvalid:
							findflag=False
						if count >= 500:
							print "failed to find a point to place the ball"
							pdb.set_trace()
					if self.verbose_ > 0:
						print ("ball at (%f, %f), dist: %f" % (pt.x(), pt.y(), md))
					xloc, yloc = pt.x_asint(), pt.y_asint()	
					pt = gm.Point(xloc, yloc)
					#determine if the ball can be placed at the chosen position or not
					isok = True
					for j in range(i):
						dist = pt.distance(allpos[j])
						if self.verbose_ > 0:
							print ("placement dist:", dist)
						isok = isok and dist > (allr[j] + r)
					if isok:
						placeflag = False
						allr.append(r)
						allpos.append(pt)
			self.world_.add_object(bdef, initPos=gm.Point(xloc, yloc))
		return allpos

	#apply intial forces on the balls
	def apply_force(self, model):
		fs = []
		for i in range(self.numBalls_):
			ballname = 'ball-%d' % i
			if self.oppForce_:
				pos1      = self.world_.get_object_position(ballname)
				ballname2 = 'ball-%d' % np.mod(i+1,2)
				pos2      = self.world_.get_object_position(ballname2)
				ff        = pos2 - pos1 
				mag       = self.fmn_ + np.floor(self.rand_.rand()*(self.fmx_ - self.fmn_))			
				ff.make_unit_norm()
				ff.scale(mag)
				fx, fy = ff.x(), ff.y()
				if self.verbose_ > 0:
					print (fx, fy)
				print ('loc 1')
			else:
				rnd1, rnd2 = self.rand_.rand(), self.rand_.rand()
				fdiff      = self.fmx_ - self.fmn_
				if self.verbose_ > 0:
					print ('loc 2 - %f, %f' % (rnd1 * fdiff, rnd2 * fdiff))
					print ('min/max - %f, %f' % (self.fmx_, self.fmn_))
				#sample magnitude
				fmag   = self.fmn_ + np.floor(rnd1 * (self.fmx_ - self.fmn_))
				#sample theta
				ftheta = rnd1 * np.pi 			
				if self.rand_.rand() > 0.5:
					ftheta = -ftheta
				fx, fy = fmag * np.cos(ftheta), fmag * np.sin(ftheta)
			if self.verbose_ > 0:
				print ('force - fx: %f, fy: %f' % (fx, fy))
			f  = gm.Point(fx, fy)
			model.apply_force(ballname, f, forceT=1.0) 
			fs.append(f)
		return model, fs


def save_nonrect_arena_val(numseq=100):
	sv = datasaver(wthick=20, isrect=False, mxforce=1e+5, wlen=300,
								 mnseqlen=10, mxseqlen=100, wtheta=[23, 38, 45, 53])
	sv.save(numseq=numseq)	

def save_nonrect_arena_train(numseq=10000, oppforce=False, numballs=1, svprefix=None, 
														 mnwlen=500, mxwlen=800, arenasz=1600, mnforce=3e+4, 
														 mxforce=8e+4):
	drname = '/data0/pulkitag/projphysics/'
	sv = datasaver(rootpath=drname, wthick=30, isrect=False, mnforce=mnforce, mxforce=mxforce, 
								 mnwlen=mnwlen, mxwlen=mxwlen, numballs=numballs,
								 mnseqlen=10, mxseqlen=200, mnballsz=25, mxballsz=25, wtheta=[30, 60],
								 arenasz=arenasz, svprefix=svprefix)
	sv.save(numseq=numseq)	


def save_rect_arena(numSeq=10000, oppForce=False, numBalls=1, svPrefix=None,
										mnForce=3e+4, mxForce=8e+4, mnWLen=300, mxWLen=550, arenaSz=700,
										mnSeqLen=10, mxSeqLen=200, randSeed=None):

	print (mnForce, mxForce)
	drName = '/data0/pulkitag/projphysics/'
	sv = DataSaver(wThick=30, isRect=True, mnForce=mnForce, mxForce=mxForce, 
								 mnWLen=mnWLen, mxWLen=mxWLen, mnSeqLen=mnSeqLen, mxSeqLen=mxSeqLen,
								 numBalls=numBalls, mnBallSz=25, mxBallSz=25, arenaSz=arenaSz,
								 oppForce=oppForce, svPrefix=svPrefix, rootPath=drName, 
								 randSeed=randSeed)
	sv.save(numSeq=numSeq)	

def save_multishape_rect_arena(numseq=1000, numballs=1, oppforce=False):
	drname = '/data1/pulkitag/projphysics/'
	sv = datasaver(rootpath=drname,numballs=numballs,wthick=20, isrect=True, mnforce=1e+4,
								 mxforce=1e+5, mnwlen=400, mxwlen=400,
								 mnseqlen=40, mxseqlen=40, mnballsz=25, mxballsz=25, oppforce=oppforce)
	sv.save(numseq=numseq)	



class CustomWorld:
	def __init__(self, numballs=1, ballsz=25, wthick=30,
							 isrect=True, wtheta=30, wlen=300,
							 arenasz=667, verbose=0, ballLocs=[gm.Point(120,120)],
							 forces=[gm.Point(5e+4, 5e+4)], **kwargs):
		self.numballs_   = numballs
		self.bsz_  = ballsz
		self.bloc_ = copy.deepcopy(ballLocs)
		self.forces_  = copy.deepcopy(forces)
		self.wthick_  = wthick
		self.wtheta_  = wtheta
		self.isrect_  = isrect
		self.wlen_    = wlen
		self.asz_     = arenasz
		self.verbose_ = verbose
		self.xSz_     = arenasz
		self.ySz_     = arenasz

	def _generate_model(self):
		#get the coordinates of the top point
		#create the world
		self.world_ = pm.World(xSz=self.asz_, ySz=self.asz_)
		#add the walls
		self.walls_   = self.add_walls()
		#add the balls
		self.add_balls()
		#create physics simulation
		self.model_   = pm.Dynamics(self.world_)

	def generate_model(self):
		self._generate_model()
		self.apply_force()

	def get_glimpse(self, ballNum=None, cropSz=128):
		self.model_.step()
		im = self.model_.generate_image()
		print (im.shape)
		position = np.zeros((2*self.numballs_,))
		velocity = np.zeros((2*self.numballs_,))
		imBalls = []
		for b in range(self.numballs_):
			ballName = 'ball-%d' % b
			ball     = self.model_.get_object(ballName)
			pos      = ball.get_position()
			vel      = ball.get_velocity()
			position[2*b] = pos.x()
			position[2*b+1] = pos.y()
			velocity[2*b] = vel.x()
			velocity[2*b+1] = vel.y()
			xMid, yMid = round(pos.x()), round(pos.y())
			if cropSz is not None:
				imBall = 255 * np.ones((cropSz, cropSz,3)).astype(np.uint8)
				x1, x2 = max(0, xMid - cropSz/2.0), min(self.xSz_, xMid + cropSz/2.0)
				y1, y2 = max(0, yMid - cropSz/2.0), min(self.ySz_, yMid + cropSz/2.0)
				print xMid, x1, x2, cropSz
				#xSz, ySz = x2 - x1, y2 - y1
				imX1 = int(cropSz/2.0  - (xMid - x1))
				imX2 = int(cropSz/2.0  + (x2 - xMid))
				imY1 = int(cropSz/2.0  - (yMid - y1))
				imY2 = int(cropSz/2.0  + (y2 - yMid))
				imBall[imY1:imY2,imX1:imX2,:] = im[int(y1):int(y2), int(x1):int(x2),0:3]
				imBalls.append(imBall)
		return imBalls, position, velocity

	def apply_force(self):
		for b in range(self.numballs_):
			ballname = 'ball-%d' % b
			self.model_.apply_force(ballname, self.forces_[b], forceT=1.0) 

	def add_walls(self, xleft=20, yleft=20, hLen=None, vLen=None, 
								fColor=pm.Color(1.0, 0.0, 0.0)):
		if hLen is None:
			hLen = self.wlen_
		if vLen is None:
			vLen = self.wlen_
		wtheta = self.wtheta_
		if type(wtheta) == list:
			th1, th2, th3 = wtheta
		else:
			th1, th2, th3 = -wtheta, wtheta, 180-wtheta
		th1, th2, th3 = fix_theta_range(th1), fix_theta_range(th2), fix_theta_range(th3)
		walls = self._create_walls(xleft, yleft, (th1, th2, th3),
							 (vLen, hLen, vLen), fColor=fColor)
		return walls		

	##
	def _create_walls(self, xleft, yleft, thetas, wlens, fColor):
		theta1, theta2, theta3 = thetas
		wlen1,  wlen2,  wlen3  = wlens
		pt1      = gm.Point(xleft, yleft)
		dir1     = gm.theta2dir(theta1)
		pt2      = pt1 + (wlen1 * dir1)
		dir2     = gm.theta2dir(theta2)
		pt3      = pt2 + (wlen2 * dir2)
		dir3     = gm.theta2dir(theta3)
		pt4      = pt3 + (wlen3 * dir3)
		pts      = [pt1, pt2, pt3, pt4]
		if self.verbose_ > 0:
			print 'Wall Points: %s, %s, %s, %s' % (pt1, pt2, pt3, pt4)
		walls    = pm.create_cage(pts, wThick = self.wthick_, fColor=fColor)	
		#get the lines within which the balls need to be added. 
		self.pts    = pts
		self.lines_ = []
		for w in walls:
			self.world_.add_object(w)
		for i in range(len(pts)):
			self.lines_.append(gm.Line(pts[i], pts[np.mod(i+1, len(pts))]))		
		return walls

	#generates and adds the required number of balls. 
	def add_balls(self):
		for i in range(self.numballs_):
			bdef = pm.BallDef(radius=self.bsz_, fColor=pm.Color(0.5, 0.5, 0.5))
			self.world_.add_object(bdef, initPos=self.bloc_[i])


def fix_theta_range(theta):
	theta = np.mod(theta, 360)
	if theta > 180:
		theta = -(360 - theta) 
	return theta
	
def stats_force():
	#datadir = '/work5/pulkitag/projphysics/trainv2-asz700_wlen300-550_nb1_bsz25-25_f3e+04-8e+04_slen10-200_wth30/'
	datadir = '/work5/pulkitag/projphysics/trainv2-asz700_wlen200-550_nb2_bsz25-25_f8e+03-5e+04_slen10-200_wth30_oppfrc'
	theta = []
	for i in range(10000):
		if np.mod(i,1000)==1:
			print(i)
		seqfolder = osp.join(datadir, 'seq%06d' % i)
		wfile     = osp.join(seqfolder, 'data.mat')
		data      = sio.loadmat(wfile)
		force     = data['force'][:,0]
		theta.append(np.arctan2(float(force[0]), float(force[1])))
	return theta

		
def delete_garbage():
	datdir = '/work5/pulkitag/projphysics/asz667_wlen300_nb1_bsz15-35_f1e+03-1e+05_slen10-100_wth2030-_wtheta60/seq%06d/'
	for i in range(6000):
		seqdir  = datdir % (i)
		imfile  = seqdir + 'im%06d.jpg'
		datfile = os.path.join(seqdir, 'data.mat')
		dat     = sio.loadmat(datfile, squeeze_me=True)
		n       = dat['position'].shape[1] 		
		for j in range(n,100):
			imname = imfile % j
			if os.path.exists(imname):
				os.remove(imName)	


def test_custom_world():
	cw = CustomWorld(wtheta=[0, 90, -180], verbose=1)
	cw.generate_model()
	plt.ion()
	for i in range(100000):
		imList, pos, vel = cw.get_glimpse(cropSz=512)
		plt.imshow(imList[0])
		plt.show()
		ip = raw_input()
		if ip=='q':
			return
	return imList, pos, vel