Raster map.py-code
From GeoMod
To run the code, once you have VPython installed, you can either.
- Download this file: raster_map.py
- or, copy and paste the following code to a file (named "raster_map.py" for example) using your favorite text editor.
from Numeric import *
from string import *
from visual import *
from random import uniform, seed
import cPickle
class color_map:
'''For contour coloring'''
def __init__(self, zmin=-10001.0, zmax=-10000.0, contour_interval=1.0,
colmin = vector(0,0,0), colmax=vector(1,1,1), color_scheme="linear"):
self.zmin = zmin
self.zmax = zmax
self.colmin = colmin
self.colmax = colmax
self.contour_interval=contour_interval
self.color_scheme = color_scheme
if self.color_scheme == "linear":
self.set_linear_colors()
def set_linear_colors(self):
self.colors = []
z = self.zmin
while z <= self.zmax:
'''interpolate color vectors'''
zval = z #(z + z+ self.contour_interval) / 2.0
if zval >= self.zmax:
self.colors.append(self.colmax)
else:
vf = ((zval - self.zmin) / float(self.zmax - self.zmin))
v1 = self.colmin + ((zval - self.zmin) / float(self.zmax - self.zmin)) * (self.colmax - self.colmin)
#v1 = vector(vf * abs(self.colmax.x-self.colmin.x), vf * abs(self.colmax.y-self.colmin.y), vf * abs(self.colmax.z-self.colmin.z))
self.colors.append(v1)
#print z, self.colors[-1]
z += self.contour_interval
def get_color(self, zval):
col_index = int((len(self.colors)-1) * (max(zval, self.zmin) - self.zmin) / float(self.zmax - self.zmin)) #make sure zval > 0
#print "self.zmin", self.zmin
#print "self.zmax", self.zmax
#print "a", zval, max(zval, self.zmin), float(self.zmax - self.zmin), col_index, self.colors[min(len(self.colors)-1, col_index)]
#col_index = col_index * len(self.colors)
#print "b", col_index
return self.colors[min(len(self.colors)-1, col_index)]
class draw_state:
def __init__(self, draw=False, visible=False):
self.draw = draw
self.visible = visible
'''A raster map created with a 2d array'''
class raster_map:
def __init__(self, data_array, xmin=0.0, ymax=0.0, dx=1, zscale=1.0): #datarr, xmin, ymax, dx, nrows, ncols):
'''data_array is a 2d array '''
self.data = data_array
(self.nrows, self.ncols) = self.data.shape
self.xmin = xmin
self.ymax = ymax
self.dx = dx
self.scale = zscale
self.color = color_map()
self.particles = []
#self.particle_trails = []
self.default = -99999
self.imin=0
self.imax=self.default
self.jmin=0
self.jmax=self.default
self.bimin = 0
self.bimax = self.ncols
self.bjmin = 0
self.bjmax = self.nrows
'''set up switches'''
self.switch = {}
self.switch["l_line_3d"] = draw_state()
self.switch["l_draw_map"] = draw_state()
self.switch["l_map_stripe"] = draw_state()
self.switch["l_block_contour"] = draw_state()
self.switch["l_draw_boundary"] = draw_state()
self.switch["l_draw_vectors"] = draw_state()
self.switch["l_surface_3d"] = draw_state()
def pickle_map(self, outfile):
'''outfile is the file handle for the output file'''
dump(self.data, outfile)
dump(self.nrows, outfile)
dump(self.ncols, outfile)
dump(self.dx, outfile)
dump(self.scale, outfile)
'''self.color'''
dump(self.color.zmin, outfile)
dump(self.color.zmax, outfile)
dump(self.color.colmin.x, outfile)
dump(self.color.colmin.y, outfile)
dump(self.color.colmin.z, outfile)
dump(self.color.colmax.x, outfile)
dump(self.color.colmax.y, outfile)
dump(self.color.colmax.z, outfile)
dump(self.color.contour_interval, outfile)
dump(self.color.color_scheme, outfile)
'''particles'''
nparticles = len(self.particles)
dump(nparticles, outfile)
for i in self.particles:
i.pickle_particle(outfile)
dump(self.default, outfile)
dump(self.imin, outfile)
dump(self.imax, outfile)
dump(self.jmin, outfile)
dump(self.jmax, outfile)
dump(self.bimin, outfile)
dump(self.bimax, outfile)
dump(self.bjmin, outfile)
dump(self.bjmax, outfile)
dump(self.switch, outfile)
'''line_3d'''
if self.switch["l_line_3d"].draw:
dump(self.radius, outfile)
if self.switch["l_map_stripe"].draw:
dump(self.nstripes, outfile)
if self.switch["l_draw_boundary"].draw:
dump(self.bound.radius, outfile)
if self.switch["l_draw_vectors"].draw:
dump(self.vscale, outfile)
if self.switch["l_surface_3d"].draw:
dump(self.smoothe, outfile)
def redraw_all(self):
'''use when redrawing from pickled file'''
def get_x(self, c):
return self.xmin + c * self.dx
def get_y(self, r):
return self.ymax - r * self.dx
def get_val(self, pos):
r, c = self.get_rc(pos)
return self.data[r,c]
def get_rc(self, pos):
r = int(((self.ymax - pos.y + self.dx/2.0) ) /self.dx )
c = int(((pos.x - self.xmin + self.dx/2.0) ) /self.dx)
return (r, c)
def inbounds(self, pos, padding=0.0):
'''determine if pos is inside the area of the raster map'''
## print "inbounds:", padding
## print "inbounds: ", pos.x, self.xmin, padding
## print "inbounds: ", pos.x, self.xmin+(self.ncols-1)*self.dx -padding
## print "inbounds: ", pos.y, self.ymax, padding
## print "inbounds: ", pos.y, self.ymax-(self.nrows-1)*self.dx+padding
## if pos.x < self.xmin+padding:
## print "out of bounds xmin"
## if pos.x > self.xmin+(self.ncols-1)*self.dx -padding:
## print "out of bounds xmax"
## if pos.y > self.ymax-padding:
## print "out of bounds ymax"
## if pos.y < self.ymax-(self.nrows-1)*self.dx+padding:
## print "out of bounds ymin"
if pos.x < self.xmin+padding or \
pos.x > self.xmin+(self.ncols-1)*self.dx -padding or \
pos.y > self.ymax-padding or \
pos.y < self.ymax-(self.nrows-1)*self.dx+padding:
#print "out of bounds"
l_in = False
else:
l_in = True
return l_in
def get_quad_nodes(self, pos):
'''returns nodes surrounding pos in the order topleft, bottomleft, topright, bottom right'''
(r,c) = self.get_rc(pos)
n = []
x = self.get_x(c)
y = self.get_y(r)
if pos.x >= x and pos.y >= y: # upper right quad
#print "ur"
n.append((r-1, c))
n.append((r,c))
n.append((r-1, c+1))
n.append((r, c+1))
elif pos.x >= x and pos.y < y: # lower right quad
#print "lr"
n.append((r,c))
n.append((r+1, c))
n.append((r, c+1))
n.append((r+1, c+1))
elif pos.x < x and pos.y < y: # upper left quad
#print "ul"
n.append((r-1, c-1))
n.append((r, c-1))
n.append((r-1, c))
n.append((r,c))
else: # lower left quad
#print "ll"
n.append((r, c-1))
n.append((r+1, c-1))
n.append((r,c))
n.append((r+1, c))
return n
def add_particle(self, pos, radius=None, color=None , l_trail=False, trail_radius=0.0, trail_color=vector(0,0.5,0.5)):
if not radius:
radius = self.dx/5.0
if not color:
color = vector(0,0,1)
pos.z = self.interpolate_linear(pos)
#print "particle position = ", pos
#test = sphere(pos=pos, color=(0,0,1), radius=self.dx)
#self.particles.append(sphere(pos=pos, radius=radius, color=color))
self.particles.append(particle(pos=pos, radius=radius, color=color,
l_trail=l_trail, trail_radius=trail_radius, trail_color=trail_color))
def delete_particle(self, particle_number):
self.particles[particle_number].delete_particle()
#last_particle = self.particles.pop(particle_number)
#if self.particles[particle_number].l_trail:
# self.particles[particle_number].delete_trail()
self.last_particle = self.particles.pop(particle_number)
def delete_all_particles(self):
if len(self.particles) > 0:
for n, i in enumerate(self.particles):
i.delete_particle()
for n, i in enumerate(self.particles):
last_particle = self.particles.pop(n)
def add_particles_at_nodes(self, step_skip = 1, l_with_trails=False, l_show_trails=False):
for i in range(1, self.nrows-1, step_skip):
for j in range(1, self.ncols-1, step_skip):
self.add_particle(pos=vector(self.dx*i, -self.dx*j), l_trail=l_with_trails)
if l_with_trails and not l_show_trails:
self.particles[-1].hide_trail()
def particles_show(self):
if len(self.particles) > 0:
for i in self.particles:
i.particle.visible = 1
def particles_hide(self):
if len(self.particles) > 0:
for i in self.particles:
i.particle.visible = 0
def draw_trails(self, radius=0.0, color=vector(0, 0.5, 0.5)):
if len(self.particles) > 0:
for i in self.particles:
i.create_trail(radius=radius, color=color)
def trails_show(self):
if len(self.particles) > 0:
for i in self.particles:
i.show_trail()
def trails_hide(self):
#print "trails hide:", len(self.particles)
#fish = chicken - goats
if len(self.particles) > 0:
for i in self.particles:
i.hide_trail()
#fish = chicken +goats
def reset_trails(self):
if len(self.particles) > 0:
for i in self.particles:
i.delete_trail()
def interpolate_linear(self, pos):
n = self.get_quad_nodes(pos)
#print n
(x1, y1, z1) = (self.get_x(n[0][1]), self.get_y(n[0][0]), self.data[n[0][0], n[0][1]])
z2 = self.data[n[1][0], n[1][1]]
z3 = self.data[n[2][0], n[2][1]]
z4 = self.data[n[3][0], n[3][1]]
wx = (pos.x - x1) / self.dx
wy = (y1 - pos.y) / self.dx
zp = (z1 * (1.0-wx) + z3*wx) * (1.0-wy) + (z2 * (1.0-wx) + z4*wx) * (wy)
return zp
def interpolate_onarray(self, pos, arry):
n = self.get_quad_nodes(pos)
(x1, y1, z1) = (self.get_x(n[0][1]), self.get_y(n[0][0]), arry[n[0][0], n[0][1]])
z2 = arry[n[1][0], n[1][1]]
z3 = arry[n[2][0], n[2][1]]
z4 = arry[n[3][0], n[3][1]]
wx = (pos.x - x1) / self.dx
wy = (y1 - pos.y) / self.dx
zp = (z1 * (1.0-wx) + z3*wx) * (1.0-wy) + (z2 * (1.0-wx) + z4*wx) * (wy)
return zp
'''Draw 3d maps as a set of blocks'''
def draw_map(self, imin=0, imax=-99999,
jmin=0, jmax=-99999,
scale = 1.0, center = 0,
bzmin = 99999, bzmax = -99999,
col_min = 99999, col_max = -99999,
shade=2, contour_interval = 10.0, color_map=None):
''' draw a specific region of the raster map'''
self.scale = scale
self.shade = shade
self.contour_interval = contour_interval
self.switch["l_draw_map"].draw = True
self.switch["l_draw_map"].visible = True
if imax == self.default:
imax = self.nrows
if jmax == self.default:
jmax = self.ncols
print "drawing map "
print "imin, imax ", imin, imax
print "jmin, jmax ", jmin, jmax
#minimum and maximum indicies of boxes
self.bimin=imin
self.bimax=imax
self.bjmin=jmin
self.bjmax=jmax
if bzmin == 99999:
self.bzmin = 99999
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
if self.data[i,j] < self.bzmin:
self.bzmin = self.data[i,j]
else:
self.bzmin = bzmin
if bzmax == -99999:
self.bzmax = -99999
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
if self.data[i,j] > self.bzmax:
self.bzmax = self.data[i,j]
else:
self.bzmax = bzmax
if col_min == 99999:
self.bcol_min = self.bzmin #- (self.bzmax-self.bzmin)*0.05
else:
self.bcol_min = col_min
if col_max == -99999:
self.bcol_max = self.bzmax
else:
self.bcol_max = col_max
if self.bcol_max == self.bcol_min:
self.bcol_max += 1
print "bzmin, bzmax ", self.bzmin, self.bzmax
print "bcol_min, bcol_max", self.bcol_min, self.bcol_max
#tdata = self.data[imin:imax,jmin:jmax]
'''Set up color scale'''
seed(1)
if color_map:
self.color = color_map
self.shade = 3
else:
if self.shade == 2:
self.color_scale = contour_colors(self.bzmin, self.bzmax, self.contour_interval)
print "shade =", self.shade
#draw in topography
#tmax = max(max(tdata))
#tmin = min(min(tdata))
self.boxes = [] #list of boxes that will be drawn
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = scale * (self.data[i,j] + self.bzmin) / 2.0
#colval = (self.data[i,j] - self.bcol_min) / (self.bcol_max - self.bcol_min) #greyscale
#print " kpos, colval =", self.data[i,j], colval
self.boxes.append(box(pos=(ipos,jpos, kpos),
length = self.dx,
height = self.dx,
width = (self.data[i,j] - self.bzmin) * scale))
#color = (colval,colval,colval)))
if color_map:
self.boxes[-1].color = self.color.get_color(self.data[i,j])
elif self.shade == 1: #put in a color shade on area
colval = 1.0 #greyscale
self.boxes[-1].color = (colval,colval,colval)
elif self.shade == 2:
self.boxes[-1].color = self.color_scale.find_color(self.data[i,j])
if center == 1:
self.scene_center(imin, imax, jmin, jmax)
def update_draw_map(self, bzmin=99999, bzmax=-99999): #draw change in level of water
ct = -1
#print self.bimin, self.bimax, self.bjmin, self.bjmax
if self.shade != 3:
'''Reset color scale'''
if bzmin == 99999:
self.bzmin = 99999
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
if self.data[i,j] < self.bzmin:
self.bzmin = self.data[i,j]
else:
self.bzmin = bzmin
if bzmax == -99999:
self.bzmax = -99999
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
if self.data[i,j] > self.bzmax:
self.bzmax = self.data[i,j]
else:
self.bzmax = bzmax
seed(1)
if self.shade == 2:
self.color_scale = contour_colors(self.bzmin, self.bzmax, self.contour_interval)
for i in range(self.bimin,self.bimax):
for j in range(self.bjmin,self.bjmax):
ct += 1
#print self.bcol_min, self.bcol_max, self.data[i,j]
self.boxes[ct].pos.z = self.scale * (self.data[i,j] + self.bzmin) / 2.0
self.boxes[ct].width = (self.data[i,j] - self.bzmin) * self.scale
if self.shade ==3:
self.boxes[ct].color = self.color.get_color(self.data[i,j])
elif self.shade == 2:
self.boxes[ct].color = self.color_scale.find_color(self.data[i,j])
else:
colval = (self.data[i,j] - self.bcol_min) / (self.bcol_max - self.bcol_min)
self.boxes[ct].color = vector(colval, colval, colval)
def draw_map_hide(self):
self.switch["l_draw_map"].visible = False
for i in boxes:
i.visible = 0
def draw_map_show(self):
self.switch["l_draw_map"].visible = True
for i in boxes:
i.visible = 1
def scene_center(self, imin, imax, jmin, jmax):
self.xc = self.xmin + (self.dx * ((jmin+jmax)/2.0))
self.yc = self.ymax - (self.dx * ((imin+imax)/2.0))
self.zc = self.get_val(vector(self.xc,self.yc,0))
print "center", self.xc, self.yc
scene.center = vector(self.xc, self.yc, self.zc)
def map_stripe(self, nstripes=5):
ct=-1
dbound = (self.bzmax - self.bzmin) / float(nstripes)
self.switch["l_map_stripe"].draw= True
self.switch["l_map_stripe"].visible= True
self.nstripes = nstripes
seed(1)
lay_cols = []
for k in range(nstripes):
lay_cols.append(vector(uniform(0.0,1.0), uniform(0.0,1.0),uniform(0.0,1.0)))
for i in range(self.bimin,self.bimax):
for j in range(self.bjmin,self.bjmax):
ct += 1
self.boxes[ct].color = color.white
for k in range(self.nstripes):
#print i, j, k, self.data[i,j], self.bzmin, (self.bzmin+(float(k+1) * dbound))
if (self.data[i,j] > (self.bzmin+(float(k) * dbound))
and self.data[i,j] <= (self.bzmin+(float(k+1) * dbound)) ):
colval = k % 2
gcol = (k+1)/nstripes
#print k, colval
self.boxes[ct].color = lay_cols[k] #(1-gcol,gcol,1)
break
def block_contour(self, contour_interval=10.0, contour_color=0):
'''contour_interval is the contour interval
contour_color = 0: Random colors for the contours
= 1: alternating red and white colours
'''
self.switch["l_block_contour"].draw = True
self.switch["l_block_contour"].visible = True
ct = -1
lay_cols = []
contour_boxes = []
zmin = int(self.bzmin/contour_interval) * contour_interval
zmax = (int(self.bzmax/contour_interval) + 1) * contour_interval
n_intervals = int((zmax - zmin) / contour_interval)
#choose colors
if contour_color == 0:
for i in range(n_intervals+1):
lay_cols.append(vector(uniform(0.0,1.0), uniform(0.0,1.0),uniform(0.0,1.0)))
elif contour_color == 1:
for i in range(n_intervals+1):
if remainder(i , 2) == 1:
lay_cols.append( vector(1,1,1) )
else:
lay_cols.append( vector(1,0,0) )
# print "zmin", zmin, "zmax", zmax, "n_intervals", n_intervals
for i in range(self.bimin,self.bimax):
for j in range(self.bjmin,self.bjmax):
ct+=1
ht = zmin
lay = -1
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
while ht < self.data[i,j]:
lay += 1
#print "lay", lay, "data", self.data[i,j], "ht", ht
if (ht + contour_interval) < self.data[i,j]: #add full block
contour_boxes.append(box(pos=(ipos,jpos, self.scale*(ht+(contour_interval/2.0))),
length = self.dx,
height = self.dx,
width = contour_interval * self.scale,
color = lay_cols[lay] ))
#print "width", contour_boxes[-1].width, "midpt", contour_boxes[-1].pos.z
else:
self.boxes[ct].width = self.scale * (self.data[i,j] - ht)
self.boxes[ct].pos.z = self.scale * (ht + ((self.data[i,j]-ht)/2.0))
self.boxes[ct].color = lay_cols[lay]
#print "width", self.boxes[ct].width, "midpt", self.boxes[ct].pos.z
ht += contour_interval
def draw_boundary(self, radius=1, center=1):
self.switch["l_draw_boundary"].draw = True
self.switch["l_draw_boundary"].visible = True
self.bound = curve(color=color.blue, radius=radius)
self.bound.append(pos=(self.xmin, self.ymax,self.data[0,0]))
self.bound.append(pos=(self.xmin+self.dx*(self.ncols-1), self.ymax,self.data[0,0]))
self.bound.append(pos=(self.xmin+self.dx*(self.ncols-1),
self.ymax-self.dx*(self.nrows-1),self.data[0,0]))
self.bound.append(pos=(self.xmin,
self.ymax-self.dx*(self.nrows-1),self.data[0,0]))
self.bound.append(pos=(self.xmin, self.ymax,self.data[0,0]))
if center == 1:
self.scene_center(self.bimin, self.bimax, self.bjmin, self.bjmax)
def line_3d(self, imin=0, imax=-99999, jmin=0, jmax=-99999,
scale=1.0, center=1, color_map=None, radius=None):
''' draw a specific region of the raster map as a wire mesh
imin - minimum column value to be drawn
imax - maximum column value
jmin - miinimum row
jmax - maximum row
scale - to add a vertical exageration
center - to center scene on area drawn if center == 1'''
self.scale = scale
default_val = -99999
self.switch["l_line_3d"].draw = True
self.switch["l_line_3d"].visible = True
if radius:
self.radius = radius
else:
self.radius = 0
if color_map:
self.color = color_map
elif not self.color:
self.color = color_map()
if imax == default_val:
imax = self.nrows
if jmax == default_val:
jmax = self.ncols
print
#tdata = self.data[imin:imax,jmin:jmax]
self.imin=imin
self.imax=imax
self.jmin=jmin
self.jmax=jmax
self.lin = []
for i in range(imin, imax):
self.lin.append(curve(color=(0,1,0), radius = self.radius))
for j in range (jmin, jmax):
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = self.data[i,j]
self.lin[-1].append(pos=(ipos,jpos,kpos*scale), color=self.color.get_color(kpos))
#print lin.pos[-1]
for j in range (jmin, jmax):
self.lin.append(curve(color=(0,1,0), radius=self.radius))
for i in range(imin, imax):
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = self.data[i,j]
self.lin[-1].append(pos=(ipos,jpos,kpos*scale), color=self.color.get_color(kpos))
#center of scene of area
#print self.xmin + jmin* self.dx, self.ymax - imax* self.dx
if center == 1:
self.scene_center(imin, imax, jmin, jmax)
## if center == 1:
## xc = self.xmin + (self.dx * ((jmin+jmax)/2))
## yc = self.ymax - (self.dx * ((imin+imax)/2))
## zc = self.get_val(vector(xc,yc,0))
## print "center", xc, yc
## scene.center = vector(xc, yc, zc)
## print "list lenght = ", len(self.lin)
def update_line_3d(self):
n = 0
#print "imin, imax", self.imin, self.imax
for i in range(self.imin, self.imax):
#lin.append(curve(color=(0,1,0)))
c = 0
#print "jmin, jmax", self.jmin, self.jmax
for j in range (self.jmin, self.jmax):
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = self.data[i,j]
#print n, c, "-", ipos, jpos, kpos
#print self.lin[n].pos[c]
self.lin[n].pos[c]=(ipos, jpos, kpos*self.scale)
self.lin[n].color[c]=self.color.get_color(kpos)
c += 1
#print lin.pos[-1]
n += 1
for j in range (self.jmin, self.jmax):
#lin.append(curve(color=(0,1,0)))
c=0
for i in range(self.imin, self.imax):
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = self.data[i,j]
#print n, c, "-", ipos, jpos, kpos
self.lin[n].pos[c]=(ipos, jpos, kpos*self.scale)
self.lin[n].color[c]=self.color.get_color(kpos)
c+=1
n += 1
def line_3d_hide(self):
self.switch["l_line_3d"].visible = False
for i in self.lin:
i.visible = 0
def line_3d_show(self):
self.switch["l_line_3d"].visible = True
for i in self.lin:
i.visible = 1
def calc_surface_vectors(self):
'''find d(data)/dx and d(data)/dy'''
self.vx = 0.0 * ones((self.nrows, self.ncols))
self.vy = 0.0 * ones((self.nrows, self.ncols))
self.vx[self.bimin+1:self.bimax-1, self.bjmin+1:self.bjmax-1] = \
-(self.data[self.bimin+1:self.bimax-1, self.bjmin+2:self.bjmax] - \
self.data[self.bimin+1:self.bimax-1, self.bjmin:self.bjmax-2] ) / (2.0*self.dx)
self.vy[self.bimin+1:self.bimax-1, self.bjmin+1:self.bjmax-1] = \
(self.data[self.bimin+2:self.bimax, self.bjmin+1:self.bjmax-1] - \
self.data[self.bimin:self.bimax-2, self.bjmin+1:self.bjmax-1] ) / (2.0*self.dx)
def draw_vectors(self, vscale=1.0, color_map=None):
self.vscale = vscale
self.switch["l_draw_vectors"].draw = True
self.switch["l_draw_vectors"].visible = True
self.calc_surface_vectors()
print self.vx
print self.vy
self.arrows = []
ct = -1
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
ct += 1
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = self.scale * self.data[i,j]
#print i,j
self.arrows.append(arrow(pos=(ipos,jpos,kpos),
axis = (self.vx[i,j]*self.vscale, self.vy[i,j]*self.vscale, 0.0)))
if not color_map:
self.arrows[-1].color = color.red
def vectors_show(self):
self.switch["l_draw_vectors"].visible = True
for i in self.arrows:
i.visible = 1
def vectors_hide(self):
self.switch["l_draw_vectors"].visible = False
for i in self.arrows:
i.visible = 0
def update_vectors(self, l_calc = True):
#print "updating vectors"
if l_calc:
self.calc_surface_vectors()
ct = -1
for i in range(self.bimin, self.bimax):
for j in range(self.bjmin, self.bjmax):
ct += 1
ipos = self.xmin + j* self.dx
jpos = self.ymax - i* self.dx
kpos = self.scale * self.data[i,j]
self.arrows[ct].pos = vector(ipos,jpos,kpos)
self.arrows[ct].axis = vector(self.vx[i,j]*self.vscale, self.vy[i,j]*self.vscale, 0.0)
def surface_3d(self, imin=0, imax=-99999, jmin=0, jmax=-99999,
scale=1.0, center=1, twoSided=True, smooth = False, color_map=None):
''' draw a specific region of the raster map as a wire mesh
imin - minimum column value to be drawn
imax - maximum column value
jmin - miinimum row
jmax - maximum row
scale - to add a vertical exageration
center - to center scene on area drawn if center == 1'''
self.scale = scale
self.switch["l_surface_3d"].draw = True
self.switch["l_surface_3d"].visible = True
if color_map:
self.color = color_map
default_val = -99999
self.smooth = smooth
if imax == default_val:
imax = self.nrows
if jmax == default_val:
jmax = self.ncols
#tdata = self.data[imin:imax,jmin:jmax]
self.imin=imin
self.imax=imax
self.jmin=jmin
self.jmax=jmax
self.surf = surface_model(twoSided=twoSided, color_map = self.color)
pt1 = zeros( (3,), Float)
pt2 = zeros( (3,), Float)
pt3 = zeros( (3,), Float)
pt4 = zeros( (3,), Float)
ct = 0
for i in range(imin, imax-1):
for j in range(jmin, jmax-1):
pt1[0] = self.xmin + j* self.dx
pt2[0] = self.xmin + (j+1)* self.dx
pt3[0] = self.xmin + (j+1)* self.dx
pt4[0] = self.xmin + j* self.dx
pt1[1] = self.ymax - i* self.dx
pt2[1] = self.ymax - i* self.dx
pt3[1] = self.ymax - (i+1)* self.dx
pt4[1] = self.ymax - (i+1)* self.dx
pt1[2] = self.data[i,j] * self.scale
pt2[2] = self.data[i,j+1] * self.scale
pt3[2] = self.data[i+1,j+1] * self.scale
pt4[2] = self.data[i+1,j] * self.scale
if (pt1[2] <> 0.0 and
pt2[2] <> 0.0 and
pt3[2] <> 0.0 and
pt4[2] <> 0.0):
#color = self.color.get_color
self.surf.FacetedPolygon(pt1, pt2, pt3, pt4)
ct +=1
#print "ct=", ct
if self.smooth == 1:
self.surf.DoSmoothShading()
if center == 1:
self.scene_center(imin, imax, jmin, jmax)
def update_surface2_3d(self):
#tmp = self.surf.model.__copy__()
tmp = self.surf.frame.objects[0].__copy__()
print tmp
##
for i in self.surf.frame.objects:
#print i
i.visible = 0
self.surface_3d()
## tmp.visible = 0
def update_surface_3d(self):
tct = 0
pt1 = zeros( (3,), Float)
pt2 = zeros( (3,), Float)
pt3 = zeros( (3,), Float)
pt4 = zeros( (3,), Float)
for i in range(self.imin, self.imax-1):
for j in range(self.jmin, self.jmax-1):
pt1[0] = self.xmin + j* self.dx
pt2[0] = self.xmin + (j+1)* self.dx
pt3[0] = self.xmin + (j+1)* self.dx
pt4[0] = self.xmin + j* self.dx
pt1[1] = self.ymax - i* self.dx
pt2[1] = self.ymax - i* self.dx
pt3[1] = self.ymax - (i+1)* self.dx
pt4[1] = self.ymax - (i+1)* self.dx
pt1[2] = self.data[i,j] * self.scale
pt2[2] = self.data[i,j+1] * self.scale
pt3[2] = self.data[i+1,j+1] * self.scale
pt4[2] = self.data[i+1,j] * self.scale
if (pt1[2] <> 0.0 and
pt2[2] <> 0.0 and
pt3[2] <> 0.0 and
pt4[2] <> 0.0):
tct = self.surf.UpdateFacetedPolygon(tct, pt1, pt2, pt3, pt4)
if self.smooth == 1:
self.surf.DoSmoothShading()
## self.surf.DoSmoothShading()
## for i in range(len(self.model.pos)):
## #print "pos", self.model.pos[i]
## lay_num = int(( self.model.pos[i][2] - zmin ) / contour_interval)
## #print "lay_num", lay_num
## self.model.color[i] = lay_cols[lay_num]
##
def ContourColor(self, contour_color = 0,
contour_interval=10.0, zmin=0.0, zmax=1000.0):
self.surf.ColorContour(contour_color,contour_interval*self.scale,
zmin*self.scale, zmax*self.scale)
def local_3x3(self, pos):
#find the local 3x3 matrix of topography around the cell centered at pos
r, c = self.get_rc(pos)
#print "r,c=", r, c
return self.data[r-1:r+2, c-1:c+2]
def get_slope(self, pos, dir):
#the slope given the positon (pos) and direction and the adjacent cells
'''WARNING: THIS IS A PROTOTYPE VERSION AND SHOULD BE USED WITH CAUTION'''
print "WARNING: THIS IS A PROTOTYPE VERSION (sub_raster_array) AND SHOULD BE USED WITH CAUTION"
posz = self.get_val(pos)
pnp = pos + (dir*self.dx)
pnelev = self.get_val(pnp)
slope = (pnelev - posz/self.scale)/(self.dx)
#print "r,c=", r, c
return slope
def sub_raster_array(self, center_pos, cell_range=1):
'''to create a sub raster using only cells within 1 cell_range of the
center cell'''
r, c = self.get_rc(center_pos)
xmin = self.get_x(c - cell_range)
ymax = self.get_y(r - cell_range)
#print center_pos
#print "x, y", self.get_x(c), self.get_y(r)
return raster_map(self.local_3x3(center_pos), xmin, ymax, self.dx)
def write_raster_map(self, outfile_name="raster_map_out.txt"):
outfile = open(outfile_name, "w")
outln = "ncols " + str(self.ncols) + "\n"
outln += "nrows " + str(self.nrows) + "\n"
outln += "xllcorner " + str(self.xmin) + "\n"
outln += "yllcorner " + str(self.ymax) + "\n"
outln += "cellsize " + str(self.dx) + "\n"
outln += "NODATA_value " + "-9999" + "\n"
outfile.write(outln)
for i in range(self.nrows):
outln = ""
for j in range(self.ncols):
outln = outln + str(self.data[i, j]) + " "
outln = outln + "\n"
outfile.write(outln)
outfile.close()
#########################
'''END OF RASTER_MAP CLASS'''
#########################
class surface_model:
def __init__(self, twoSided= True, color_map=None):
self.frame = frame()
self.model = faces(frame=self.frame)
self.twoSided = twoSided # add every face twice with opposite normals
if color_map:
self.color = color_map
else:
self.color = color_map()
def FacetedTriangle(self, v1, v2, v3, color=color.white):
"""Add a triangle to the model, apply faceted shading automatically"""
v1 = vector(v1)
v2 = vector(v2)
v3 = vector(v3)
try:
normal = norm( cross(v2-v1, v3-v1) )
except:
normal = vector(0,0,0)
#print "self.twoSided = ", self.twoSided
for v in (v1,v3,v2):
self.model.append( pos=v, color=self.color.get_color(v[2]), normal=normal )
#print "initial position = ", self.model.pos[-1]
#print v
if self.twoSided:
for v in (v1,v2,v3):
self.model.append( pos=v, color=self.color.get_color(v[2]), normal=-normal )
#print v
#print "done"
#rate(1)
def FacetedPolygon(self, *v):
"""Appends a planar polygon of any number of vertices to the model,
applying faceted shading automatically."""
for t in range(len(v)-2):
self.FacetedTriangle( v[0], v[t+1], v[t+2], color )
def UpdateFacetedTriangle(self, tct, v1, v2, v3, color=color.white):
"""Add a triangle to the model, apply faceted shading automatically"""
v1 = vector(v1)
v2 = vector(v2)
v3 = vector(v3)
try:
normal = norm( cross(v2-v1, v3-v1) )
except:
normal = vector(0,0,0)
for v in (v1,v3,v2):
#print "tct, old_pos, new_pos", tct, self.model.pos[tct], v
self.model.pos[tct] = v
self.model.normal[tct] = normal
self.model.color[tct] = self.color.get_color(v[2])
#print tct, v
tct += 1
if self.twoSided:
for v in (v1,v2,v3):
self.model.pos[tct] = v
self.model.normal[tct] = normal
self.model.color[tct] = self.color.get_color(v[2])
#print tct, v
tct += 1
#print "done, normal =", normal
return tct
def UpdateFacetedPolygon(self, tct, *v):
"""Appends a planar polygon of any number of vertices to the model,
applying faceted shading automatically."""
for t in range(len(v)-2):
tct = self.UpdateFacetedTriangle( tct, v[0], v[t+1], v[t+2])
return tct
def DoSmoothShading(self):
"""Change a faceted model to smooth shaded, by averaging normals at
coinciding vertices.
This is a very slow and simple smooth shading
implementation which has to figure out the connectivity of the
model and does not attempt to detect sharp edges.
It attempts to work even in two-sided mode where there are two
opposite normals at each vertex. It may fail somehow in pathological
cases. """
pos = self.model.pos
normal = self.model.normal
vertex_map = {} # vertex position -> vertex normal
vertex_map_backface = {}
for i in range( len(pos) ):
tp = tuple(pos[i])
old_normal = vertex_map.get( tp, (0,0,0) )
if dot(old_normal, normal[i]) >= 0:
vertex_map[tp] = normal[i] + old_normal
else:
vertex_map_backface[tp] = normal[i] + vertex_map_backface.get(tp, (0,0,0))
for i in range( len(pos) ):
tp = tuple(pos[i])
if dot(vertex_map[tp], normal[i]) >= 0:
normal[i] = vertex_map[tp] and norm( vertex_map[ tp ] )
else:
normal[i] = vertex_map_backface[tp] and norm(vertex_map_backface[tp] )
def DrawNormal(self, scale):
pos = self.model.pos
normal = self.model.normal
for i in range(len(pos)):
arrow(pos=pos[i], axis=normal[i]*scale)
def ColorContour(self, contour_color = 0,
contour_interval=10.0, zmin=0.0, zmax=1000.0):
ct = -1
lay_cols = []
zmin = int(zmin/contour_interval) * contour_interval
zmax = (int(zmax/contour_interval) + 1) * contour_interval
n_intervals = int((zmax - zmin) / contour_interval)
print "n_intervals", n_intervals
print "zmin, zmax", zmin, zmax
#choose colors
if contour_color == 0 or contour_color == 2: #random colors
for i in range(n_intervals+1):
lay_cols.append(vector(uniform(0.0,1.0), uniform(0.0,1.0),uniform(0.0,1.0)))
elif contour_color == 1: #alternating red and white
for i in range(n_intervals+1):
if remainder(i , 2) == 1:
lay_cols.append( vector(1,1,1) )
else:
lay_cols.append( vector(1,0,0) )
if contour_color == 0 or contour_color == 1: #sharp coloring
for i in range(len(self.model.pos)):
#print "pos", self.model.pos[i]
lay_num = int(( self.model.pos[i][2] - zmin ) / contour_interval)
#print "lay_num", lay_num
self.model.color[i] = lay_cols[lay_num]
elif contour_color == 2: #smoothe coloring
for i in range(len(self.model.pos)):
lay_a = int(( self.model.pos[i][2] - zmin ) / contour_interval)
wt_a = ( (lay_a * contour_interval + zmin) - self.model.pos[i][2] ) / contour_interval
if (lay_a == 1 and wt_a >=0.0) or (lay_a == n_intervals and wt_a <= 0.0):
out_col = lay_cols[lay_a]
elif wt_a >= 0.0:
out_col = ( lay_cols[lay_a] * wt_a ) + ( lay_cols[lay_a - 1] * (1 - wt_a) )
else:
out_col = (- lay_cols[lay_a] * wt_a ) + ( lay_cols[lay_a + 1] * (1 + wt_a) )
self.model.color[i] = out_col
def import_new_raster(self,rasterfile):
(self.data, self.xmin, self.ymax, self.dx) = read_arcgis_ascii(rasterfile)
#return raster_map(arr, xmin, ymax, dx)
def read_arcgis_ascii(rasterfile):
'''Read in an ArcGIS ASCII raster file'''
infile = open(rasterfile,"r")
ct = 0
line = infile.readline()
while line:
ct += 1
l = split(strip(line))
if ct == 1:
ncols = int(l[1])
if ct == 2:
nrows = int(l[1])
if ct == 3:
xmin = float(l[1])
if ct == 4:
ymin = float(l[1])
if ct == 5:
dx = float(l[1])
if ct == 6:
nodata = l[1]
print "nodata", nodata
if ct >= 7:
if ct == 7:
topo = zeros((nrows*ncols,), Float)
t_ct = -1
for i in l:
t_ct += 1
if i == nodata: #set no-data values to zero
topo[t_ct] = 0.0
else:
topo[t_ct] = float(i)
line = infile.readline()
infile.close
ymax = ymin + (dx * (nrows-1)) # top left corner
'''datarr is a 2d array '''
#self.xmin = xmin
#self.ymax = ymax
#self.dx = dx
#self.nrows = nrows
#self.ncols = ncols
#self.data = resize(topo, (nrows, ncols))
return (resize(topo, (nrows, ncols)), xmin, ymax, dx)
#return raster_map(resize(topo, (nrows, ncols)), xmin, ymax, dx)
def raster_import_ryan(rasterfile, dx=1.0):
'''THIS FUNCTION READS IN X, Y, DATA text
Code to add if using Ryans ASCII conversion (reminder: dx is defaulted to 1)'''
infile = open(rasterfile,"r")
old_long = []
old_lat = []
old_data = []
line = infile.readline()
l = splitfields(strip(line),",")
old_long.append(float(l[0]))
old_lat.append(float(l[1]))
old_data.append(float(l[2]))
ct = 0
#data[ct] = float(l[2])
ncols = 1
nrows = 1
while line:
l = splitfields(strip(line),",")
#print float(l[1]), float(l[2])
old_long.append(float(l[0]))
old_lat.append(float(l[1]))
old_data.append(float(l[2]))
#data[ct+1] = float(l[2])
line = infile.readline()
ct += 1
if (old_long[-1] < old_long[-2]):
nrows += 1
ncols = ct/nrows
data = zeros((nrows*ncols,), Float)
print "length = ", len(old_data), nrows*ncols
for i in range (nrows*ncols):
data[i] = old_data[i]
print "rows, columns, number of items ", nrows, ncols, ct
xmin = old_long[0]
ymax = old_lat[0]
return raster_map(resize(data, (nrows,ncols)), xmin, ymax, dx)
def unpickle_map(infile):
'''MUST MATCH THE pickle_map METHOD IN THE raster_map CLASS'''
'''infile is the file handle for the input file'''
r = raster_map( load(infile))
r.nrows = load( infile)
r.ncols = load(infile)
r.dx = load(infile)
r.scale = load(infile)
'''self.color'''
zmin = load(infile)
zmax = load(infile)
colmin_x = load(infile)
colmin_y = load(infile)
colmin_z = load(infile)
colmax_x = load(infile)
colmax_y = load(infile)
colmax_z = load(infile)
contour_interval = load(infile)
color_scheme = load(infile)
r.color = color_map(zmin=zmin, zmax=zmax, contour_interval=contour_interval,
colmin=vector(colmin_x, colmin_y, colmin_z),
colmax=vector(colmax_x, colmax_y, colmax_z),
color_scheme = color_scheme)
'''particles'''
r.particles = []
nparticles = load(infile)
for i in range(nparticles):
r.particles.append(unpickle_particle(infile))
r.default = load(infile)
r.imin = load(infile)
r.imax = load(infile)
r.jmin = load(infile)
r.jmax = load(infile)
r.bimin = load( infile)
r.bimax = load(infile)
r.bjmin = load(infile)
r.bjmax = load(infile)
r.switch = load(infile)
'''line_3d'''
if r.switch["l_line_3d"].draw:
r.line_3d(radius=load(infile))
if r.switch["l_map_stripe"].draw:
r.map_stripe(nstripes=load(infile))
if r.switch["l_draw_boundary"].draw:
r.draw_boundary(radius=load(infile))
if r.switch["l_draw_vectors"].draw:
r.draw_vectors(vscale=load(infile))
if r.switch["l_surface_3d"].draw:
r.surface_3d(smoothe=load(infile))
return r
class contour_colors:
'''Produces color scale for contouring'''
def __init__(self, zmin, zmax, dz = 10.0, color_scheme = "random"):
'''To set up the colors for contour levels given:
- the contour interval of dz
- the starting point of zmin
- the maximum level of zmax
Data is put into the list lay_cols'''
self.min = zmin
self.dz = dz
self.color_scheme = color_scheme
self.colors = []
z = zmin
#print "contour_colors -> z, zmax", z, zmax
if z == zmax:
self.colors.append(vector(1,1,1))
else:
while z < zmax:
if color_scheme == "random":
self.colors.append(vector(uniform(0.0,1.0), uniform(0.0,1.0),uniform(0.0,1.0)))
z += dz
def find_color(self, val):
col_index = int((val - self.min) / self.dz)
## print col_index
## print self.colors
## print self.colors[col_index]
return self.colors[col_index]
class vector_field:
'''Takes two coincident raster_map's (uraster and vraster) and produces a 2d vector field
uses the dimensions from the first raster map entered for plotting,
the first raster map MUST have already drawn a map (draw_map)'''
def __init__(self, uraster, vraster,
scale = 1.0, offset = 0.0):
self.scale = scale
self.offset = offset
self.uraster = uraster
self.vraster = vraster
'''The vector mask array can be used to draw only selected areas of a set of vectors.'''
vector_mask = resize(zeros((self.uraster.nrows*self.uraster.ncols,), Int ), (self.uraster.nrows, self.uraster.ncols))
def draw_vectors(self,
imin=0, imax=-99999,
jmin=0, jmax=-99999):
'''define indices for drawing'''
self.imin = imin
self.jmin = jmin
if imax == -99999:
self.imax = self.uraster.nrows
else:
self.imax = imax
if jmax == -99999:
self.jmax = self.uraster.ncols
else:
self.jmax = jmax
print "i, j indicies"
print self.imin, self.imax
print self.jmin, self.jmax
self.arrows = []
ct = -1
for i in range(self.imin, self.imax):
for j in range(self.jmin, self.jmax):
#print i,j
ct += 1
self.arrows.append(arrow(pos=self.uraster.boxes[ct].pos,
axis = (self.uraster.data[i,j]*self.scale, self.vraster.data[i,j]*self.scale, 0.0+self.offset)))
def draw_povray_vectors(self, export_directory, ini_base_file, pov_base_file, normal_vector = 15,
imin=0, imax=-99999,
jmin=0, jmax=-99999):
f = open(ini_base_file,"r")
ini_f = f.read()
f.close()
f = open(pov_base_file,"r")
pov_f = f.read()
f.close
#print ini_f
# print ini_f.find("<outfile>")
#print ini_f.replace("<outfile>", "CHICKEN")
print "xxccx"
## line = f.readline()
## print "xxxx", line
## while line:
## print " xxx", line
## print line.find("<outfile>")
## line = f.readline()
## pov_file = open(pov_base_file, "r").readline()
## pov_out_file = open("tmp.pov", "w")
## pov_f = pov_file.read()
##
'''define indices for drawing'''
self.imin = imin
self.jmin = jmin
if imax == -99999:
self.imax = self.uraster.nrows
else:
self.imax = imax
if jmax == -99999:
self.jmax = self.uraster.ncols
else:
self.jmax = jmax
print "i, j indicies"
print self.imin, self.imax
print self.jmin, self.jmax
old_lat = -9999
self.arrows = []
ct = -1
for i in range(self.imin, self.imax):
for j in range(self.jmin, self.jmax):
#print i,j
ct += 1
u = self.uraster.data[i,j]
v = self.vraster.data[i,j]
lat = 90 - (i*2.5)
lng = j * 2.5
print i, j, lat, lng, u, v
lat = "%+03.1F" % lat
lng = "%+03.1F" % lng
if lat <> old_lat:
out_dir = export_directory + "/lat" + lat
subprocess.Popen([r"mkdir", out_dir], stdout=subprocess.PIPE).communicate()[0]
old_lat = lat
#jucies
ini_out_file = open("tmp.ini", "w")
pov_out_file = open("tmp.pov", "w")
#outfile_name = "%+03.1F%+03.1F.png" % (self.uraster.boxes[ct].pos.x, self.uraster.boxes[ct].pos.y)
outfile_name = lat + lng + ".png"
outfile_name = out_dir + "/" + outfile_name
print outfile_name
#juices
ini = ini_f.replace("<outfile>", outfile_name)
#ini_f = ini_f.replace("<oldfile>", outfile_name)
#print ini
#print ini_out_file
ini_out_file.write(ini)
ini_out_file.close()
out_angle = get_angle_360(u, v)
out_angle = "%3.2F" % out_angle
#print "out_angle = ", out_angle
pov = pov_f.replace("<rotangle>", out_angle)
vect_magnitude = hypot(u, v) / normal_vector
vect_magnitude = "%3.5F" % vect_magnitude
pov = pov.replace("<vector_magnitude>", vect_magnitude)
pov_out_file.write(pov)
pov_out_file.close()
subprocess.Popen([r"/sw/bin/povray", "tmp.ini"], stdout=subprocess.PIPE).communicate()[0]
#juices
'''modify arrow-basic.pov file'''
## self.arrows.append(arrow(pos=self.uraster.boxes[ct].pos,
## axis = (self.uraster.data[i,j]*self.scale, self.vraster.data[i,j]*self.scale, 0.0+self.offset)))
class layer_raster:
'''holds top and bottom elevations of a layer as raster_map's
takes two inputs, a top and bottom of a layer'''
def __init__(self, top, bot):
self.top = top
self.bot = bot
self.color = color_map()
def draw_layer(self, color_map=None, center=1):
'''color_map is a class defined in this file'''
## if color == "white":
## self.color = color_map()
if color_map:
self.color = color_map
print self.color.colors
self.boxes = [] #list of boxes that will be drawn
for i in range(self.top.bimin, self.top.bimax):
for j in range(self.top.bjmin, self.top.bjmax):
ipos = self.top.xmin + j* self.top.dx
jpos = self.top.ymax - i* self.top.dx
boxtop = self.top.scale * self.top.data[i,j]
boxbot = self.top.scale * self.bot.data[i,j]
kpos = (boxtop + boxbot) / 2.0
#colval = (self.data[i,j] - self.bcol_min) / (self.bcol_max - self.bcol_min) #greyscale
#print " kpos, colval =", self.data[i,j], colval
self.boxes.append(box(pos=(ipos,jpos, kpos),
length = self.top.dx,
height = self.top.dx,
width = (boxtop - boxbot)*self.top.scale,
color = self.color.get_color(self.top.data[i,j])))
#print self.boxes[-1].color
if center == 1:
self.top.scene_center(self.top.bimin, self.top.bimax, self.top.bjmin, self.top.bjmax)
def update_layer(self):
ct = -1
#print self.bimin, self.bimax, self.bjmin, self.bjmax
#print self.shade
for i in range(self.top.bimin,self.top.bimax):
for j in range(self.top.bjmin,self.top.bjmax):
ct += 1
boxtop = self.top.scale * self.top.data[i,j]
boxbot = self.top.scale * self.bot.data[i,j]
kpos = (boxtop + boxbot) / 2.0
#print self.bcol_min, self.bcol_max, self.data[i,j]
self.boxes[ct].pos.z = kpos
self.boxes[ct].width = (boxtop - boxbot)*self.top.scale
self.boxes[ct].color = self.color.get_color(self.top.data[i,j])
def layer_show(self):
for i in self.boxes:
i.visible = 1
def layer_hide(self):
for i in self.boxes:
i.visible = 0
def draw_layer_surfaces(self, color_map_top = None, color_map_bot=None):
if color_map_top:
self.color_top = color_map_top
else:
self.color_top = color_map()
if color_map_bot:
self.color_bot = color_map_bot
else:
self.color_bot = self.color_top
self.top.surface_3d(color_map=self.color_top)
self.bot.surface_3d(color_map=self.color_bot)
def update_layer_surfaces(self):
self.top.update_surface_3d()
self.bot.update_surface_3d()
def unpickle_particle(infile):
'''infile is the file handle'''
px = load(infile)
py = load(infile)
pz = load(infile)
ox = load(infile)
oy = load(infile)
oz = load(infile)
cx = load(infile)
cy = load(infile)
cz = load(infile)
radius = load(infile)
l_trail = load(infile)
draw_limit = load(infile)
pvis = load(infile)
if l_trail:
tradius = load(infile)
tcolor = load(infile)
tpos = load(infile)
tvis = load(infile)
p = particle(pos=vector(px, py, pz), radius=radius, color=vector(cx, cy,cz),
l_trail=l_trail, draw_limit =draw_limit)
p.old_pos = vector(ox, oy, oz)
p.particle.visible = pvis
if p.l_trail:
p.trail.pos = tpos
p.trail.radius=tradius
#print "tcolor", len(tcolor), tcolor
p.trail.color = tcolor
p.trail.visible = tvis
return p
class particle:
def __init__(self, pos, radius, color, l_trail=False, draw_limit=0.0,
trail_radius=0.0, trail_color=vector(0,0.5,0.5)):
'''draw_limit is the distance between drawing points on a particle trail'''
self.pos = pos
self.old_pos = pos
self.color = color
self.particle = sphere(pos=self.pos, radius=radius, color=self.color)
self.l_trail = l_trail
self.draw_limit = draw_limit
if self.l_trail:
self.create_trail(radius=trail_radius, color=trail_color)
def pickle_particle(self, outfile):
'''outfile is the file handle'''
dump(self.pos.x, outfile)
dump(self.pos.y, outfile)
dump(self.pos.z, outfile)
dump(self.old_pos.x, outfile)
dump(self.old_pos.y, outfile)
dump(self.old_pos.z, outfile)
dump(self.color.x, outfile)
dump(self.color.y, outfile)
dump(self.color.z, outfile)
dump(self.particle.radius, outfile)
dump(self.l_trail, outfile)
dump(self.draw_limit, outfile)
dump(self.particle.visible, outfile)
if self.l_trail:
dump(self.trail.radius, outfile)
dump(self.trail.color, outfile)
dump(self.trail.pos, outfile)
dump(self.trail.visible, outfile)
'''must match with function unpickle_particle'''
def move_particle(self, new_pos):
self.particle.pos = self.pos = new_pos
if self.l_trail:
self.add_to_trail()
def create_trail(self, radius=0.0, color=vector(0, 0.5, 0.5)):
self.l_trail = True
#print "color=", color, radius
self.trail = curve( radius=radius)
self.trail.append(pos=self.particle.pos, color=color)
def add_to_trail(self, draw_limit=None):
#print "magnitude=", self.old_pos, self.particle.pos, mag(self.old_pos, self.particle.pos)
if draw_limit:
self.draw_limit = draw_limit
if mag(self.old_pos - self.particle.pos) >= self.draw_limit:
self.trail.append(pos=self.particle.pos)
self.old_pos = self.particle.pos
def show_trail(self):
self.trail.visible = 1
def hide_trail(self):
if self.l_trail:
self.trail.visible = 0
def delete_trail(self):
if self.l_trail:
self.l_trail = False
self.hide_trail()
self.trail.pos = []
def delete_particle(self):
self.particle.visible = 0
if self.l_trail:
self.trail.visible = 0
def get_angle_360(u,v):
ang = degrees(atan2(v, u))
if ang < 0.0:
ang = ang + 360
return ang
