spring.py

用于OMNeT++的模板

#!/usr/bin/python

from curves import readfile
from types import *
import sys,math,os

sys.path.append(os.path.join(os.getcwd(),"build/lib.linux-i686-"+".".join([str(x) for x in sys.version_info[0:2]])))
try:
	from gauss import gaussian
except:
	print "Run 'python setup.py build' first!"
	sys.exit(1)
	
import cPickle as pickle

def frange(start,stop,jump=1):
	x = start
	while x<stop:
		yield x
		x += jump

def gen_squares(out):
	box = list(out["limits"])
	squares = []
	print out["limits"]
	x,y = box[0:2]
	while box[2]-box[0]!=box[3]-box[1]:
		h = box[3]-box[1]
		w = box[2]-box[0]
		if w>h:
			squares.append({"size":h,"x":x,"y":y})
			x += h
			box[0]+=h
		else:
			squares.append({"size":w,"x":x,"y":y})
			y += w
			box[1]+=w
	squares.append({"size":box[2]-box[0],"x":x,"y":y})
	return squares
	
def draw_squares(screen,square,scaling=5):
	white = pygame.color.Color("white")
	for s in squares:
		pygame.draw.rect(screen,white,(s["x"]*scaling,s["y"]*scaling,s["size"]*scaling,s["size"]*scaling),1)

def old_squares(limits,step=0.5):
	limits = [int(x) for x in limits]
	return reduce (lambda x,y:x+y,[[{"size":step,"x":x,"y":y} for x in frange(limits[0],limits[2],step)] for y in frange(limits[1],limits[3],step)])

def prob_squares(squares,anchors,total,mult):
	for s in squares:
		s["p"] = 1
	total = 1
	print "mult",mult
	for a in anchors:
		#print "a",a
		newtotal = 0
		if type(a[0])==ListType:
			r = a[2]
			todo = a[0:2]
			#print "todo",todo
		else:
			r = a[3]
			todo = [a[0:3]]
			
		for s in squares:
			x = s["x"] + (s["size"]/2.0)
			y = s["y"] + (s["size"]/2.0)
		
			val = 0
			for t in todo:
				dist = math.sqrt((t[0]-x)**2+(t[1]-y)**2)
				val += t[2]*gaussian((dist - (s["size"]/2.0))*mult,(dist + (s["size"]/2.0))*mult,r,14*0.2)

			if type(a[0])==ListType:
				val /= 4.0
			
			s["p"] *= val/total
			newtotal += s["p"]
		#print "newtotal",newtotal
		total = newtotal
	for s in squares:
		s["p"] /= total
	#print "prob total",total
	
	return squares

def old_prob(big, small,step):
	for b in big:
		set = [s["p"] for s in small if s["x"]>=b["x"] and s["x"] < b["x"]+b["size"] and s["y"]>=b["y"] and s["y"] < b["y"]+b["size"]]
		print b,len(set)
		if len(set)!=(b["size"]*(1/step))**2:
			raise Exception,"Should be "+str((b["size"]*(1/step))**2)+" not "+str(len(set))
		b["p2"] = reduce(lambda x,y:x+y,set)

if not os.path.exists(sys.argv[1]+".pickle"):
	out = readfile(sys.argv[1])
else:
	out = pickle.load(open(sys.argv[1]+".pickle"))
total = reduce(lambda x,y:x+y,[reduce(lambda x,y:x+y,z) for z in out["data"]])
#print out
print "anchors",out["ref"]
squares = gen_squares(out)
#print squares
#print "calced",
#calced = prob_squares(squares,out["ref"],1,0.871)
calced = prob_squares(squares,out["ref"],1,1.0)
#print "old",
step = 0.05
if not out.has_key("refined") or not out["refined"].has_key(step):
	old = prob_squares(old_squares(out["limits"],step),out["ref"],1,1.0)
	#print "old",old
	if not out.has_key("refined"):
		out["refined"] = {}

	out["refined"][step] = old
	pickle.dump(out,open(sys.argv[1]+".pickle",'w'),pickle.HIGHEST_PROTOCOL)
else:
	old = out["refined"][step]

old_prob(calced,old,step)
print calced
diff = 0
for dat in calced:
	diff += math.fabs(dat["p"]-dat["p2"])
	print "dat",dat["size"],dat["p"],dat["p2"],math.fabs(dat["p"]-dat["p2"])/dat["p2"]
	
h = out["limits"][3]-out["limits"][1]
w = out["limits"][2]-out["limits"][0]

print "improvement",((w*h)-len(squares))/(w*h)*100,"%"
print "total",total
print "diff",diff,str(diff*100)+"%","avg",str(diff*100/len(calced))+"%"

sys.exit(0)

import pygame

pygame.init()
screen = pygame.display.set_mode((800, 800))

clock = pygame.time.Clock()
while 1:
	clock.tick(10)
	pygame.display.flip()