Raster map.py

From GeoMod

Jump to: navigation, search

from string import *
from visual import *
from random import uniform, seed
import cPickle
try:
    from Numeric import *
except:
    print "numpy"

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", nodata=-9999.0):
        self.zmin = zmin
        self.zmax = zmax
        self.colmin = colmin
        self.colmax = colmax
        self.contour_interval=contour_interval
        self.color_scheme = color_scheme
        self.nodata=nodata
        
        if self.color_scheme == "linear":
            self.set_linear_colors()
        if self.color_scheme == "rainbow":
            self.set_rainbow_colors()

    def set_rainbow_colors(self, color_map_file="gradient_map.txt"):
            cfile = open(color_map_file)
            l = cfile.readline()
            self.colors = []
            self.contour_interval = 1.0
            for i in range(256):
                l = cfile.readline()
                (r,g,b) = l.split("(")[1].split(")")[0].split(",")
                #print r, g, b
                self.colors.append(vector(float(r)/(256*256),float(g)/(256*256),float(b)/(256*256)))
            #print self.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
        if zval <> self.nodata:
            col = self.colors[min(len(self.colors)-1, col_index)]
        else:
            #print "nodata check"
            col = scene.background
        #print zval, col
        return col


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, nodata=-9999.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.nodata = nodata
        self.color = color_map(nodata=self.nodata)

        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 min_val(self):
        # determine the minimum value of the data array
        minval = 9e50
        ct = 0
        for i in self.data.flat:
            if i <> self.nodata and i < minval:
                minval = i
            ct += 1
        return minval

    def max_val(self):
        # determine the maximum value of the data array
        maxval = -9e50
        ct = 0
        for i in self.data.flat:
            if i <> self.nodata and i > maxval:
                maxval = i
            ct += 1
        return maxval

    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 = 1,
                 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

        # find minimum value in array
        self.minval = self.min_val()

        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 and self.data[i,j] <> self.nodata:
                        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 and self.data[i,j] <> self.nodata:
                        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)

        if center == 1:
            self.scene_center(imin, imax, jmin, jmax)

        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
                adj_elev = self.adjusted_data_val(self.data[i,j])
                kpos = scale * (adj_elev + 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 = (adj_elev - self.bzmin) * scale))
                                      #color = (colval,colval,colval)))
                if color_map:
                    self.boxes[-1].color = self.color.get_color(adj_elev)
                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(adj_elev)
                    
                    

    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 and self.data[i,j] <> self.nodata:
                            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 and self.data[i,j] <> self.nodata:
                            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]
                adj_elev = self.adjusted_data_val(self.data[i,j])
                self.boxes[ct].pos.z = self.scale * (adj_elev + self.bzmin) / 2.0 
                self.boxes[ct].width = (adj_elev - self.bzmin) * self.scale
                if self.shade ==3:
                    self.boxes[ct].color = self.color.get_color(adj_elev)
                elif self.shade == 2:
                    self.boxes[ct].color = self.color_scale.find_color(adj_elev)
                else:
                    colval = (adj_elev - self.bcol_min) / (self.bcol_max - self.bcol_min)
                    self.boxes[ct].color = vector(colval, colval, colval)

    '''get the list index for the closest box given a position (based on x,y position)'''
    def get_box_number(self, pos):
        (r, c) = self.get_rc(pos)
        return  r * self.ncols + c

    '''get reference to closest box given the position of the box'''
    def get_box_from_pos(self, pos):
        (r, c) = self.get_rc(pos)
        return  self.boxes[r * self.ncols + c]
                 
    '''get reference to closest box given the row and column (r,c) of the box'''
    def get_box_from_rc(self, r,c):
        return  self.boxes[r * self.ncols + c]
                 
    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, zcenter=True):
        self.xc = self.xmin + (self.dx * ((jmin+jmax)/2.0))
        self.yc = self.ymax - (self.dx * ((imin+imax)/2.0))
        if zcenter:
            self.zc = self.get_val(vector(self.xc,self.yc,0)) * self.scale
        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

        # find minimum value in array
        self.minval = self.min_val()

        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
        
        #center of scene of area
        if center == 1:
            self.scene_center(imin, imax, jmin, jmax)

        #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.adjusted_data_val(self.data[i,j])
                self.lin[-1].append(pos=(ipos,jpos,kpos*scale), color=self.color.get_color(self.data[i,j]))
                #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.adjusted_data_val(self.data[i,j])
                self.lin[-1].append(pos=(ipos,jpos,kpos*scale), color=self.color.get_color(self.data[i,j]))

##        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 adjusted_data_val(self, dataval):
        if dataval == self.nodata:
            dataval = self.minval
        return dataval

    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.adjusted_data_val(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(self.data[i,j])
                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.adjusted_data_val(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(self.data[i,j])
                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
        
        # find minimum value in array
        self.minval = self.min_val()

        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

        if center == 1:
            self.scene_center(imin, imax, jmin, 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.adjusted_data_val(self.data[i,j]) * self.scale
                pt2[2] = self.adjusted_data_val(self.data[i,j+1]) * self.scale
                pt3[2] = self.adjusted_data_val(self.data[i+1,j+1]) * self.scale
                pt4[2] = self.adjusted_data_val(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()

    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.adjusted_data_val(self.data[i,j]) * self.scale
                pt2[2] = self.adjusted_data_val(self.data[i,j+1]) * self.scale
                pt3[2] = self.adjusted_data_val(self.data[i+1,j+1]) * self.scale
                pt4[2] = self.adjusted_data_val(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()

    def lowRC(self, r, c):
        '''returns the row and column of the lowest cell adjacent to the cell defined by r,c'''
        if (r == 0) or r == self.nrows-1 or c == 0 or c == self.ncols-1:
            (rlow, clow) = self.lowRCbc(r,c)
        else:
            (rlow, clow) = self.lowRCi(r,c)
        return (rlow, clow)
            
        
    def lowRCi(self, r, c):
        '''returns the row and column of the lowest cell adjacent to the cell defined by r,c
           this only works for internal cells'''
        ind = argmin(self.data[r-1:r+2, c-1:c+2])
        #print r, c, ind, self.data[r-1:r+2, c-1:c+2].shape
        #print self.data[r-1:r+2, c-1:c+2]
        (rlow, clow) = unravel_index(ind, (3,3))
        return (r-1+rlow, c-1+clow)

        #return ind

    def lowRCbc(self, r, c):
        '''this variant of lowRC returns the index of the lowest cell adjacent to the boundary cell
            defined by r,c '''
        if r == 0 and c == 0:
            opts = ((1, 2), (2,1), (2,2))
        elif r == 0 and c == self.ncols-1:
            opts = ( (1,0), (2,0), (2,1))
        elif r == self.nrows-1 and c == 0:
            opts = ((0,1), (0,2), (1,2))
        elif r == self.nrows-1 and c == self.ncols-1:
            opts = ((0,0), (0,1), (1,0))
        elif r == 0:
            opts = ((1,0), (1,2), (2,0), (2,1), (2,2)) 
        elif r == self.nrows-1:
            opts = ((0,0), (0,1), (0,2), (1,0), (1,2))
        elif c == 0:
            opts = ((0,1), (0,2), (1,2), (2,1), (2,2))
        elif c == self.ncols-1:
            opts = ((0,0), (0,1), (1,0), (2,0), (2,1))
        else:
            print "whoops, not on boundary (in lowRCbc)"

        low_val = self.data[r, c]
        (rl, cl) = (1, 1) #the center position in the 3x3 array
        #print "cell =", r, c
        for i in opts:
            #print i, r-1+i[0], c-1+i[1]
            if self.data[r-1+i[0], c-1+i[1]] < low_val:
                low_val = self.data[r-1+i[0], c-1+i[1]]
                (rl, cl) = i
        (rlow, clow) = (r-1+rl, c-1+cl)
        return (rlow, clow)

    def build_catchment(self, l_draw_trace=0):
        '''construct an array with the amount of cells contributing flow to each cell'''
        self.catchment = ones((self.nrows,self.ncols), int)
        self.trace_number = zeros((self.nrows,self.ncols), int)
        for r in range(self.nrows):
            for c in range(self.ncols):
                if self.data[r,c] <> self.nodata:
                    flow_trace = [(r,c),]
                    (rlow, clow) = (r,c)
                    while 0 < rlow < self.nrows and 0 < clow< self.ncols:
                        # find the next lowest cell
                        (rlow, clow) = self.lowRC(rlow,clow)
                        # check if we've already been to this cell or if out of bounds
                        for i in flow_trace:
                            if rlow == i[0] and clow == i[1]:
                                rlow = -1 # set r=-1 as a sign to skip out of loops
                            if self.data[rlow, clow] == self.nodata:
                                rlow = -1
                        if rlow <> -1:
                            flow_trace.append((rlow,clow))
                    # draw trace
                    if (l_draw_trace==1):
                        #print "trace", r, c
                        trace_pos = []
                        for i in flow_trace:
                            trace_pos.append( (self.get_x(i[1]), self.get_y(i[0]), self.data[i[0], i[1]] ) )
                            #print (self.get_x(i[1]), self.get_y(i[0]), self.data[i[0], i[1]] ) 
                        l = curve(pos=trace_pos, radius=1)
                    #mouse_pause(scene)

    def write_ImageMagick_text_image(self, fname, color_map="greyscale"):
        '''Write array to a greyscale file that can be read by ImageMagick (convert) and
            converted to other formats using the command line "convert" command.'''
        outf = open(fname, "w")

        if color_map == "greyscale":
            '''Write header'''
            outl = "# ImageMagick pixel enumeration: {0},{1},255,RGB \n".format(self.ncols, self.nrows)
            outf.write(outl)
            for i in range(self.nrows):
                for j in range(self.ncols):
                    outl = "{0},{1}: ({2:3d},{2:3d},{2:3d}) \n".format(j,i,int(self.data[i,j]))
                    outf.write(outl)

        if color_map == "greyscale16":
            '''Write header'''
            outl = "# ImageMagick pixel enumeration: {0},{1},65535,RGB \n".format(self.ncols, self.nrows)
            outf.write(outl)
            for i in range(self.nrows):
                for j in range(self.ncols):
                    outl = "{0},{1}: ({2:3d},{2:3d},{2:3d}) \n".format(j,i,int(self.data[i,j]))
                    outf.write(outl)

        if color_map == "rainbow":
            '''read in color map'''
            colors = []
            cfile = open("gradient_map.txt")
            l = cfile.readline()
            for i in range(256):
                l = cfile.readline()
                colors.append( l.split("(")[1].split(")")[0].split(","))
                #print colors[-1]

            
            '''Write header'''
            outl = "# ImageMagick pixel enumeration: {0},{1},65535,RGB \n".format(self.ncols, self.nrows)
            outf.write(outl)
            '''write file'''
            for i in range(self.nrows):
                for j in range(self.ncols):
                    elev = int(self.data[i,j])
                    outl = "{0},{1}: ({2:3d},{3:3d},{4:3d}) \n".format(j,i , int(colors[elev][0]), int(colors[elev][1]), int(colors[elev][2]) )
                    outf.write(outl)

        if color_map == "rainbow16":
            '''read in color map'''
            colors = []
            cfile = open("gradient_map16.txt")
            l = cfile.readline()
            for i in range(256*256):
                l = cfile.readline()
                colors.append( l.split("(")[1].split(")")[0].split(","))
                #print colors[-1]

            
            '''Write header'''
            outl = "# ImageMagick pixel enumeration: {0},{1},65535,RGB \n".format(self.ncols, self.nrows)
            outf.write(outl)
            '''write file'''
            for i in range(self.nrows):
                for j in range(self.ncols):
                    elev = int(self.data[i,j])
                    outl = "{0},{1}: ({2:3d},{3:3d},{4:3d}) \n".format(j,i , int(colors[elev][0]), int(colors[elev][1]), int(colors[elev][2]) )
                    outf.write(outl)


            
        outf.close()


    
        
        
##     def build_catchment(self, l_draw_trace=0):
##        '''construct an array with the amount of cells contributing flow to each cell'''
##        trace_endpoint_number = 10000000
##        self.catchment = ones((self.nrows,self.ncols), int)
##        # trace_r and trace_c are the indicies of the first trace that hit the specified cell
##        self.trace_r = zeros((self.nrows,self.ncols), int)
##        self.trace_c = zeros((self.nrows,self.ncols), int)
##        self.trace_num = ones((self.nrows,self.ncols), int) * -1 # let -1 be the default value
##        self.traces = []
##        trace_number = -1
##        for r in range(self.nrows):
##            for c in range(self.ncols):
##                catchment_number = 1 # the number of cells contributing to this cell
##                if self.data[r,c] <> self.nodata:
##                    trace_number += 1
##                    self.traces.append(trace((r,c)))
##                    #self.trace[-1] = [(r,c),]
##                    #flow_trace = [(r,c),]
##                    (rlow, clow) = (r,c)
##                    (self.trace_r[rlow,clow], self.trace_c[rlow,clow]) = (rlow, clow)
##                    self.trace_num[rlow,clow] = trace_number
##                    while 0 < rlow < self.nrows and 0 < clow< self.ncols:
##                        # FIND THE NEXT LOWEST CELL
##                        (rlow, clow) = self.lowRC(rlow,clow)
##                        # check if we've gone out of bounds
##                        if self.data[rlow, clow] == self.nodata:
##                            rlow = -1
##                        # check if we've already been to this cell
##                        elif self.trace_num[rlow, clow] == trace_number:
##                            rlow = -1
##                        # check if we're in a new cell
##                        elif self.trace_num[rlow, clow] == -1:
##                            catchment_number += 1
##                        # if the existing trace is not the same as this trace then follow the existing trace downhill
##                        elif self.trace_num[rlow, clow] <> trace_number:
##                            for i in self.trace[self.trace_num[rlow, clow]]:
##                                
##                            # mark the current trace as ending at the existing trace number
##                            self.trace[-1].append((trace_endpoint_number+trace_number,-1))
##                            # the first number in the tuple refers to the new endpoint
##                            # the second number in the tuple refers to the location on the 
##                            
##                                    
##                                
##                        else:
##                            
##
##                        
##                        for i in self.trace[-1]:
##                            if self.data[rlow, clow] == self.nodata:
##                                rlow = -1
##                            elif self.trace_num[rlow, clow] <> -1:
##                                # follow the existing trace while adding the catchement value to all cells in main stream
##                                rlow = -1 # set r=-1 as a sign to skip out of loops
##                        # check if we've reached an existing trace
##                        
##                        if rlow <> -1:
##                            self.trace[-1].append((rlow,clow))
##                            self.trace_num[rlow,clow] = trace_number
##                    # draw trace
##                    if (l_draw_trace==1):
##                        #print "trace", r, c
##                        trace_pos = []
##                        for i in flow_trace:
##                            trace_pos.append( (self.get_x(i[1]), self.get_y(i[0]), self.data[i[0], i[1]] ) )
##                            #print (self.get_x(i[1]), self.get_y(i[0]), self.data[i[0], i[1]] ) 
##                        l = curve(pos=trace_pos, radius=1)
##                    #mouse_pause(scene)

##   def trace_flow(self, r,c):
##        flow_trace = ((r,c),)
##        (rlow, clow) = (r,c)
##        while 0 < rlow < self.nrows and 0 < clow< self.ncols:
##            # find the next lowest cell
##            (rlow, clow) = self.lowRC(r,c)
##            # check if we've already been to this cell
##            for i in flow_trace:
##                if rlow == i[0] and clow == i[1]:
##                    rlow = -1 # set r=-1 as a sign to skip out of loops
##            if rlow <> -1:
##                flow_trace.append((rlow,clow))
##        return flow_trace
        

            
                

'''Testing code'''    
##    def catchment(self,tt):
##        self.c_out = zeros((self.ny,self.nx), int)
##        self.r_out = zeros((self.nx,self.ny), int)
##        self.cr_out = zeros((self.nx,self.ny), float)
##        for i in range(1,self.nx-1):
##            for j in range(1,self.ny-1):
##                self.LowTrace(i,j)
##                
##    '''The following function is also the subroutine for the function 'catchment' it sends the array for 3x3
##    surronding the given cell to function 'LowRC'''
##    def LowTrace(self,r,c): 
##        self.PosR = r
##        self.PosC = c
##        (LowRow, LowCol, Low) = self.LowRC()
##        self.cr_out[r,c] = self.ElevS[LowRow, LowCol]
##        self.c_out[r,c] = LowCol
##        self.r_out[r,c] = LowRow
##
##    '''The following function is the subroutine for the function 'catchment' it calculates the
##        lowest value in that array of pixels '''
##    def LowRC(self): 
##        Low = self.ElevS[self.PosR,self.PosC]
##        self.PosRL = self.PosR
##        self.PosCL = self.PosC
##        temp =0
##        for i in range(self.PosR-1,self.PosR+2):
##            for j in range(self.PosC-1,self.PosC+2):
##                if Low > self.ElevS[i,j]:
##                     Low = self.ElevS[i,j]
##                     self.PosRL = i
##                     self.PosCL = j
##                   
##        return (self.PosRL, self.PosCL, Low)        
        

''' End testing code'''

#########################
    
'''END OF RASTER_MAP CLASS'''

#########################

class trace:
    def __init__(self, start_point=(-1,-1)):
        self.list = []
        self.list.append(start_point)
        self.endtrace = -1 # set to the existing trace_number if this trace ends at another trace.
        self.merge_pos = -1 # set the merge position within an existing trace

    def find_merge_position(r, c):
        # this function finds the position in the list of the tuple (r,c)
        l_check = 0
        for i in range(len(self.list)):
            if r == self.list[i][0] and c == self.list[i][1]:
                l_check = 1
                break
        if l_check == 0:
            self.passionfruit()
        else:
            return i
            
            
        


            
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, self.nodata) = read_arcgis_ascii(rasterfile)
    #return raster_map(arr, xmin, ymax, dx)

def raster_import(rasterfile):
	(data, xmin,  ymax, dx, nodata) = read_arcgis_ascii(rasterfile)
	return raster_map(data, xmin=xmin, ymax=ymax, dx=dx, nodata=nodata)
                    
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
        try: #old Numeric
        	l = split(strip(line))
        except: 
        	l = line.split()
        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 = float(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
                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, nodata)
    #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


def mouse_pause(scene_x):
    l_stop = 0
    while l_stop == 0:
        m1 = scene_x.mouse.getevent()
        if m1.release:
            l_stop = 1


def rescale_array(data, minval=99999, maxval=-99999, minout=0.0, maxout=255.0):
    if minval == 99999:
        minval = data.flatten().min()
    if maxval == -99999:
        maxval = data.flatten().max()
    print "rescale: minval, maxval", minval, maxval
    data = minout + ((data -minval)*maxout/(maxval-minval))
    #data = data - minval
    return data


Personal tools